use std::collections::HashMap;
use anyhow::Result;
use bytes::Bytes;
use th_rs::define::*;
const ES_META_SAMPLE_RATE: &str = "sample_rate";
pub(crate) const ES_G711_DF_SAMPLE_RATE: i32 = 8000;
type Patch = HashMap<String, String>;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[repr(u8)]
pub enum FrameType {
I = 0,
P,
B,
}
#[repr(u8)]
#[derive(Clone, Debug, PartialEq, Eq, Copy)]
pub enum VideoCodec {
Unkown = 0,
H264,
H265,
Mpeg4,
Other,
Svac,
Svac3,
}
#[repr(u8)]
#[derive(Clone, Debug, PartialEq, Eq, Copy)]
pub enum AudioCodec {
Unkown = 0,
G711A,
Aac,
Mp3,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Video {
pub frame_type: FrameType,
pub codec: VideoCodec,
pub height: i32,
pub width: i32,
pub fps: i32,
pub time_stamp: u64,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Audio {
pub codec: AudioCodec,
pub channels: i32,
pub bits: i32,
pub sample_rate: i32,
}
impl Audio {
pub fn patch_updata(&mut self, patch: Option<&Patch>) -> Result<()> {
let Some(patch) = patch else {
return Ok(());
};
for (k, v) in patch {
if k == ES_META_SAMPLE_RATE {
self.sample_rate = v.parse()?;
}
}
Ok(())
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct RtpFull {
pub stream_type: i32,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum EsInfo {
Video(Video),
Audio(Audio),
RtpFull(RtpFull),
Com, MpegPs,
MpegTs,
Flv,
Fmp4,
Private,
}
impl EsInfo {
pub fn stream_type(&self) -> u32 {
(match self {
EsInfo::Video(_) => TH_DATATYPE_TH_DT_VIDEO,
EsInfo::Audio(_) => TH_DATATYPE_TH_DT_AUDIO,
EsInfo::RtpFull(_) => TH_DATATYPE_TH_DT_RTP_FULL,
EsInfo::Com => TH_DATATYPE_TH_DT_COM,
EsInfo::MpegPs => TH_DATATYPE_TH_DT_PS,
EsInfo::MpegTs => TH_DATATYPE_TH_DT_TS,
EsInfo::Flv => TH_DATATYPE_TH_DT_FLV,
EsInfo::Fmp4 => TH_DATATYPE_TH_DT_FMP4,
EsInfo::Private => TH_DATATYPE_TH_DT_PRIV1,
}) as _
}
pub fn update_g711(&mut self) -> Option<&mut Audio> {
match self {
EsInfo::Audio(audio) if audio.codec == AudioCodec::G711A => Some(audio),
_ => None,
}
}
pub fn into_th_esstream_info(self) -> TH_ESStreamInfo {
let mut info = TH_ESStreamInfo::default();
match self {
EsInfo::Video(video) => {
info.streamType = TH_DATATYPE_TH_DT_VIDEO as _;
info.vCodecType = video.codec as _;
info.frameType = video.frame_type as _;
info.timestamp = video.time_stamp as _;
info.frameRate = video.fps as _;
info.video_height = video.height;
info.video_width = video.width;
},
EsInfo::Audio(audio) => {
info.streamType = TH_DATATYPE_TH_DT_AUDIO as _;
info.aCodecType = audio.codec as _;
info.aSampleRate = audio.sample_rate as _;
info.aChannels = audio.channels as _;
info.aSampleBits = audio.bits as _;
},
EsInfo::RtpFull(_) => {
info.streamType = TH_DATATYPE_TH_DT_RTP_FULL as _;
},
EsInfo::Com => {
info.streamType = TH_DATATYPE_TH_DT_COM as _;
},
EsInfo::MpegPs => {
info.streamType = TH_DATATYPE_TH_DT_PS as _;
},
EsInfo::MpegTs => {
info.streamType = TH_DATATYPE_TH_DT_TS as _;
},
EsInfo::Flv => {
info.streamType = TH_DATATYPE_TH_DT_FLV as _;
},
EsInfo::Fmp4 => {
info.streamType = TH_DATATYPE_TH_DT_FMP4 as _;
},
EsInfo::Private => {
info.streamType = TH_DATATYPE_TH_DT_PRIV1 as _;
},
}
info
}
}
impl From<TH_ESStreamInfo> for EsInfo {
fn from(info: TH_ESStreamInfo) -> Self {
match info.streamType {
0 => EsInfo::Video(Video {
frame_type: {
match info.frameType {
0 => FrameType::I,
1 => FrameType::P,
2 => FrameType::B,
_ => FrameType::P,
}
},
codec: {
match info.vCodecType {
1 => VideoCodec::H264,
2 => VideoCodec::H265,
3 => VideoCodec::Svac,
4 => VideoCodec::Mpeg4,
5 => VideoCodec::Other,
6 => VideoCodec::Svac3,
_ => VideoCodec::Unkown,
}
},
height: info.video_height,
width: info.video_width,
fps: info.frameRate as _,
time_stamp: info.timestamp,
}),
1 => EsInfo::Audio(Audio {
codec: match info.aCodecType {
1 => AudioCodec::G711A,
2 => AudioCodec::Aac,
3 => AudioCodec::Mp3,
_ => AudioCodec::Unkown,
},
bits: info.aSampleBits,
channels: info.aChannels,
sample_rate: info.aSampleRate,
}),
_ => EsInfo::Private,
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct PacketEs {
payload: Bytes,
info: EsInfo,
}
impl PacketEs {
pub fn new(payload: impl Into<Bytes>, info: EsInfo) -> Self {
Self {
payload: payload.into(),
info,
}
}
pub fn new_com(payload: impl Into<Bytes>) -> Self {
let info = EsInfo::Com;
Self {
payload: payload.into(),
info,
}
}
pub fn into_splite(self) -> (Bytes, EsInfo) {
(self.payload, self.info)
}
pub fn info(&self) -> &EsInfo {
&self.info
}
pub fn is_empty(&self) -> bool {
self.payload.is_empty()
}
pub fn is_com(&self) -> bool {
matches!(
self,
PacketEs {
payload: _,
info: EsInfo::Com
}
)
}
pub fn len(&self) -> usize {
self.payload.len()
}
pub fn into_g711a(self) -> Option<(Bytes, Audio)> {
match self.info {
EsInfo::Audio(audio) => {
return (matches!(audio, Audio { codec, .. } if codec == AudioCodec::G711A))
.then(|| (self.payload, audio));
},
_ => {
return None;
},
}
}
pub fn payload(&self) -> Bytes {
self.payload.clone()
}
pub fn is_audio_packet(&self) -> bool {
matches!(
self,
PacketEs {
payload: _,
info: EsInfo::Audio(..)
}
)
}
pub fn is_video_packet(&self) -> bool {
matches!(
self,
PacketEs {
payload: _,
info: EsInfo::Video(..)
}
)
}
pub fn is_video_key_packet(&self) -> bool {
if let PacketEs {
payload: _,
info: EsInfo::Video(Video { frame_type, .. }),
} = self
{
return frame_type == &FrameType::I;
}
return false;
}
pub fn is_mpegts(&self) -> bool {
matches!(
self,
PacketEs {
payload: _,
info: EsInfo::MpegTs
}
)
}
pub fn into_video(self) -> (Video, Bytes) {
if let PacketEs {
payload,
info: EsInfo::Video(v),
} = self
{
(v, payload)
} else {
panic!("Not a Video")
}
}
pub fn into_audio(self) -> (Audio, Bytes) {
if let PacketEs {
payload,
info: EsInfo::Audio(a),
} = self
{
(a, payload)
} else {
panic!("Not a Audio")
}
}
}