Fix M2TS/MKV roundtrip: codec_private in FMKV header, Annex B conversion

Bug 1: M2TS roundtrip dropped frames — TsMuxer converts length-prefixed
NALs to Annex B, prepends VPS/SPS/PPS from HEVCDecoderConfigurationRecord.

Bug 2: MKV remux lost codec_private — MkvStream.codec_private() now returns
data from EBML header.

FMKV header carries codec_private (base64) per video stream for lossless
M2TS roundtrip.

src/mux/m2ts.rs

@@ -35,6 +35,8 @@ pub struct M2tsStream {
     pending_frames: std::collections::VecDeque<crate::pes::PesFrame>,
     pid_to_track: Vec<(u16, usize)>,
     pes_eof: bool,
+    /// Codec private data per stream (from FMKV header).
+    stored_codec_privates: Vec<Option<Vec<u8>>>,
 }
 
 impl M2tsStream {
@@ -53,6 +55,7 @@ impl M2tsStream {
             pending_frames: std::collections::VecDeque::new(),
             pid_to_track: Vec::new(),
             pes_eof: false,
+            stored_codec_privates: Vec::new(),
         }
     }
 
@@ -91,6 +94,7 @@ impl M2tsStream {
         if let Ok(Some(m)) = meta::read_header(&mut reader) {
             let header_end = reader.stream_position()?;
             let content_size = file_size.saturating_sub(header_end);
+            let codec_privates = m.codec_privates();
             let title = m.to_title();
             let (pids, parsers, pid_to_track) = Self::setup_pes(&title.streams);
             return Ok(Self {
@@ -105,6 +109,7 @@ impl M2tsStream {
                 pending_frames: std::collections::VecDeque::new(),
                 pid_to_track,
                 pes_eof: false,
+                stored_codec_privates: codec_privates,
             });
         }
 
@@ -144,6 +149,7 @@ impl M2tsStream {
             pending_frames: std::collections::VecDeque::new(),
             pid_to_track,
             pes_eof: false,
+            stored_codec_privates: Vec::new(),
         })
     }
 }
@@ -197,6 +203,11 @@ impl crate::pes::Stream for M2tsStream {
     fn info(&self) -> &crate::disc::DiscTitle { &self.disc_title }
 
     fn codec_private(&self, track: usize) -> Option<Vec<u8>> {
+        // First check stored codec_privates from FMKV header
+        if let Some(Some(cp)) = self.stored_codec_privates.get(track) {
+            return Some(cp.clone());
+        }
+        // Fall back to parser-extracted codec_private
         let pid = self.pid_to_track.iter()
             .find(|(_, idx)| *idx == track)
             .map(|(pid, _)| *pid)?;

src/mux/meta.rs

@@ -46,6 +46,9 @@ pub enum MetaStream {
         label: String,
         #[serde(default)]
         secondary: bool,
+        /// Base64-encoded codec initialization data (HEVCDecoderConfigurationRecord, etc.)
+        #[serde(default, skip_serializing_if = "Option::is_none")]
+        codec_private: Option<String>,
     },
     #[serde(rename = "audio")]
     Audio {
@@ -74,6 +77,20 @@ pub enum MetaStream {
 }
 
 impl M2tsMeta {
+    /// Build metadata from a disc Title with optional codec_private per stream.
+    pub fn from_title_with_privates(title: &DiscTitle, codec_privates: &[Option<Vec<u8>>]) -> Self {
+        let mut meta = Self::from_title(title);
+        for (i, s) in meta.streams.iter_mut().enumerate() {
+            if let MetaStream::Video { codec_private, .. } = s {
+                if let Some(Some(cp)) = codec_privates.get(i) {
+                    use base64::Engine;
+                    *codec_private = Some(base64::engine::general_purpose::STANDARD.encode(cp));
+                }
+            }
+        }
+        meta
+    }
+
     /// Build metadata from a disc Title.
     pub fn from_title(title: &DiscTitle) -> Self {
         let streams = title
@@ -88,6 +105,7 @@ impl M2tsMeta {
                     hdr: v.hdr.id().into(),
                     label: v.label.clone(),
                     secondary: v.secondary,
+                    codec_private: None,
                 },
                 Stream::Audio(a) => MetaStream::Audio {
                     pid: a.pid,
@@ -129,6 +147,7 @@ impl M2tsMeta {
                     hdr,
                     label,
                     secondary,
+                    codec_private: _,
                 } => Stream::Video(VideoStream {
                     pid: *pid,
                     codec: codec.parse().unwrap(),
@@ -183,6 +202,21 @@ impl M2tsMeta {
             content_format: crate::disc::ContentFormat::BdTs,
         }
     }
+
+    /// Extract codec_private data per stream (from FMKV header).
+    /// Returns a Vec matching stream order — None for streams without codec_private.
+    pub fn codec_privates(&self) -> Vec<Option<Vec<u8>>> {
+        self.streams.iter().map(|s| {
+            if let MetaStream::Video { codec_private: Some(ref b64), .. } = s {
+{
+                    use base64::Engine;
+                    base64::engine::general_purpose::STANDARD.decode(b64).ok()
+                }
+            } else {
+                None
+            }
+        }).collect()
+    }
 }
 
 /// Write the metadata header to a writer. Padded to 192-byte boundary.

src/mux/mkvstream.rs

@@ -207,6 +207,22 @@ impl crate::pes::Stream for MkvStream {
     fn finish(&mut self) -> io::Result<()> { Ok(()) }
 
     fn info(&self) -> &crate::disc::DiscTitle { &self.disc_title }
+
+    fn codec_private(&self, track: usize) -> Option<Vec<u8>> {
+        let track_num = (track + 1) as u16; // MKV tracks are 1-based
+        if let Mode::Read(ref rs) = self.mode {
+            rs.codec_privates
+                .iter()
+                .find(|(tn, _)| *tn == track_num)
+                .map(|(_, data)| data.clone())
+        } else {
+            None
+        }
+    }
+
+    fn headers_ready(&self) -> bool {
+        true // MKV has all headers upfront in the EBML header
+    }
 }
 
 impl IOStream for MkvStream {

src/mux/pesout.rs

@@ -13,17 +13,28 @@ pub struct M2tsOutputStream {
 }
 
 impl M2tsOutputStream {
-    pub fn create(path: &str, title: &DiscTitle) -> io::Result<Self> {
+    pub fn create(
+        path: &str,
+        title: &DiscTitle,
+        codec_privates: &[Option<Vec<u8>>],
+    ) -> io::Result<Self> {
         let file = std::fs::File::create(path)
             .map_err(|e| io::Error::new(e.kind(), format!("m2ts://{}: {}", path, e)))?;
         let mut writer = io::BufWriter::with_capacity(4 * 1024 * 1024, file);
-        // Write FMKV metadata header so M2tsStream::open can read it back
+        // Write FMKV metadata header with codec_privates so M2tsStream::open can read them back
         if !title.streams.is_empty() {
-            let m = super::meta::M2tsMeta::from_title(title);
+            let m = super::meta::M2tsMeta::from_title_with_privates(title, codec_privates);
             super::meta::write_header(&mut writer, &m)?;
         }
         let pids = extract_pids(title);
-        Ok(Self { muxer: TsMuxer::new(writer, &pids), title: title.clone() })
+        let mut muxer = TsMuxer::new(writer, &pids);
+        // Pass codec_privates to TsMuxer for Annex B parameter set injection
+        for (i, cp) in codec_privates.iter().enumerate() {
+            if let Some(data) = cp {
+                muxer.set_codec_private(i, data.clone());
+            }
+        }
+        Ok(Self { muxer, title: title.clone() })
     }
 }
 

src/mux/resolve.rs

@@ -384,7 +384,7 @@ pub fn output(
         }
         StreamUrl::M2ts { ref path } => {
             validate_file_path(path, "m2ts")?;
-            Ok(Box::new(super::pesout::M2tsOutputStream::create(&path.to_string_lossy(), title)?))
+            Ok(Box::new(super::pesout::M2tsOutputStream::create(&path.to_string_lossy(), title, codec_privates)?))
         }
         StreamUrl::Network { ref addr } => {
             validate_network_addr(addr)?;

src/mux/tsmux.rs

@@ -13,19 +13,33 @@ pub struct TsMuxer<W: Write> {
     writer: W,
     pids: Vec<u16>,
     continuity: Vec<u8>,        // per-PID continuity counter (0-15)
+    codec_privates: Vec<Option<Vec<u8>>>, // per-track codec_private (for video parameter sets)
+    params_written: Vec<bool>,  // per-track: have we written parameter sets?
 }
 
 impl<W: Write> TsMuxer<W> {
     pub fn new(writer: W, pids: &[u16]) -> Self {
-        let continuity = vec![0u8; pids.len()];
+        let n = pids.len();
         Self {
             writer,
             pids: pids.to_vec(),
-            continuity,
+            continuity: vec![0u8; n],
+            codec_privates: vec![None; n],
+            params_written: vec![false; n],
+        }
+    }
+
+    /// Set codec_private data for a track. Used to prepend VPS/SPS/PPS
+    /// as Annex B NALs before the first keyframe in the transport stream.
+    pub fn set_codec_private(&mut self, track: usize, data: Vec<u8>) {
+        if track < self.codec_privates.len() {
+            self.codec_privates[track] = Some(data);
         }
     }
 
     /// Write a PES frame as BD-TS packets.
+    /// Video frame data is expected as length-prefixed NALUs (MKV/PES format)
+    /// and is converted to Annex B for transport stream.
     pub fn write_frame(
         &mut self,
         track: usize,
@@ -36,11 +50,31 @@ impl<W: Write> TsMuxer<W> {
             return Ok(()); // unknown track, skip
         }
         let pid = self.pids[track];
+        let is_video = (0x1011..=0x101F).contains(&pid);
+
+        // For video: convert length-prefixed NALUs to Annex B (start codes)
+        // On first keyframe, prepend parameter sets from codec_private
+        let es_data = if is_video && !data.is_empty() {
+            let mut annex_b = Vec::new();
+            // Prepend codec_private parameter sets on first keyframe
+            if !self.params_written[track] {
+                if let Some(ref cp) = self.codec_privates[track] {
+                    if let Some(params) = hvcc_to_annex_b(cp) {
+                        annex_b.extend_from_slice(&params);
+                        self.params_written[track] = true;
+                    }
+                }
+            }
+            annex_b.extend_from_slice(&length_prefixed_to_annex_b(data));
+            annex_b
+        } else {
+            data.to_vec()
+        };
 
         // Build PES packet: header + data
         let pts_90k = (pts_ns * 9 / 100_000) as u64;
-        let pes_header = build_pes_header(pid, pts_90k, data.len());
+        let pes_header = build_pes_header(pid, pts_90k, es_data.len());
-        let pes_packet = [&pes_header[..], data].concat();
+        let pes_packet = [&pes_header[..], &es_data[..]].concat();
 
         // Split into TS packets
         let mut offset = 0;
@@ -152,3 +186,63 @@ fn build_pes_header(pid: u16, pts_90k: u64, data_len: usize) -> Vec<u8> {
 
     header
 }
+
+/// Extract NAL arrays from HEVCDecoderConfigurationRecord and convert to Annex B.
+/// Returns VPS + SPS + PPS as Annex B NAL units (00 00 00 01 + NAL).
+fn hvcc_to_annex_b(hvcc: &[u8]) -> Option<Vec<u8>> {
+    // HEVCDecoderConfigurationRecord: 22 bytes header, then NAL arrays
+    if hvcc.len() < 23 {
+        return None;
+    }
+    let num_arrays = hvcc[22] as usize;
+    let mut out = Vec::new();
+    let mut offset = 23;
+
+    for _ in 0..num_arrays {
+        if offset + 3 > hvcc.len() { break; }
+        // array: 1 byte (completeness + NAL type), 2 bytes (numNalus)
+        let _nal_type = hvcc[offset] & 0x3F;
+        let num_nalus = u16::from_be_bytes([hvcc[offset + 1], hvcc[offset + 2]]) as usize;
+        offset += 3;
+
+        for _ in 0..num_nalus {
+            if offset + 2 > hvcc.len() { break; }
+            let nal_len = u16::from_be_bytes([hvcc[offset], hvcc[offset + 1]]) as usize;
+            offset += 2;
+            if offset + nal_len > hvcc.len() { break; }
+            out.extend_from_slice(&[0x00, 0x00, 0x00, 0x01]);
+            out.extend_from_slice(&hvcc[offset..offset + nal_len]);
+            offset += nal_len;
+        }
+    }
+
+    if out.is_empty() { None } else { Some(out) }
+}
+
+/// Convert length-prefixed NALUs (4-byte BE length + NAL) to Annex B
+/// (00 00 00 01 + NAL). Used for video elementary streams in TS.
+fn length_prefixed_to_annex_b(data: &[u8]) -> Vec<u8> {
+    let mut out = Vec::with_capacity(data.len());
+    let mut offset = 0;
+    while offset + 4 <= data.len() {
+        let len = u32::from_be_bytes([
+            data[offset],
+            data[offset + 1],
+            data[offset + 2],
+            data[offset + 3],
+        ]) as usize;
+        offset += 4;
+        if offset + len > data.len() {
+            break;
+        }
+        out.extend_from_slice(&[0x00, 0x00, 0x00, 0x01]);
+        out.extend_from_slice(&data[offset..offset + len]);
+        offset += len;
+    }
+    // If data doesn't look like length-prefixed NALs (no valid parse),
+    // return original data unchanged — it may already be Annex B.
+    if out.is_empty() && !data.is_empty() {
+        return data.to_vec();
+    }
+    out
+}