use crate::Result;
use crate::ts::payload::{Bytes, Null, Pat, Pes, Pmt};
use crate::ts::{AdaptationField, Pid, TsHeader, TsPacket, TsPayload};
use std::collections::HashMap;
use std::io::Read;
pub trait ReadTsPacket {
fn read_ts_packet(&mut self) -> Result<Option<TsPacket>>;
}
#[derive(Debug)]
pub struct TsPacketReader<R> {
stream: R,
pids: HashMap<Pid, PidKind>,
}
impl<R: Read> TsPacketReader<R> {
pub fn new(stream: R) -> Self {
TsPacketReader {
stream,
pids: HashMap::new(),
}
}
pub fn stream(&self) -> &R {
&self.stream
}
pub fn into_stream(self) -> R {
self.stream
}
}
impl<R: Read> ReadTsPacket for TsPacketReader<R> {
fn read_ts_packet(&mut self) -> Result<Option<TsPacket>> {
let mut reader = self.stream.by_ref().take(TsPacket::SIZE as u64);
let mut peek = [0; 1];
let eos = reader.read(&mut peek)? == 0;
if eos {
return Ok(None);
}
let (header, adaptation_field_control, payload_unit_start_indicator) =
TsHeader::read_from(peek.chain(&mut reader))?;
let adaptation_field = if adaptation_field_control.has_adaptation_field() {
AdaptationField::read_from(&mut reader)?
} else {
None
};
let payload = if adaptation_field_control.has_payload() {
let payload = match header.pid.as_u16() {
Pid::PAT => {
let pat = Pat::read_from(&mut reader)?;
for pa in &pat.table {
self.pids.insert(pa.program_map_pid, PidKind::Pmt);
}
TsPayload::Pat(pat)
}
Pid::NULL => {
let null = Null::read_from(&mut reader)?;
TsPayload::Null(null)
}
0x01..=0x1F | 0x1FFB => {
let bytes = Bytes::read_from(&mut reader)?;
TsPayload::Raw(bytes)
}
_ => {
let kind = self.pids.get(&header.pid).cloned().ok_or_else(|| {
crate::Error::invalid_input(format!("Unknown PID: header={:?}", header))
})?;
match kind {
PidKind::Pmt => {
let pmt = Pmt::read_from(&mut reader)?;
for es in &pmt.es_info {
self.pids.insert(es.elementary_pid, PidKind::Pes);
}
TsPayload::Pmt(pmt)
}
PidKind::Pes => {
if payload_unit_start_indicator {
let pes = Pes::read_from(&mut reader)?;
TsPayload::PesStart(pes)
} else {
let bytes = Bytes::read_from(&mut reader)?;
TsPayload::PesContinuation(bytes)
}
}
}
}
};
Some(payload)
} else {
None
};
if reader.limit() != 0 {
return Err(crate::Error::invalid_input(
"Unexpected remaining data in TS packet",
));
}
Ok(Some(TsPacket {
header,
adaptation_field,
payload,
}))
}
}
#[derive(Debug, Clone)]
enum PidKind {
Pmt,
Pes,
}