//! M2tsStream — BD transport stream write sink.
//!
//! Write: prepends FMKV metadata header, then muxes PES frames into
//! BD-TS. The read direction lives on the pipeline highway —
//! `m2ts://` URLs route through
//! [`super::resolve::input`] → `build_m2ts_pipeline` →
//! [`super::pipelined_stream::PipelinedPesStream`], so this type is
//! write-only.

use super::meta;
use crate::disc::{DiscTitle, Stream as DiscStream};
use std::io::{self, Write};

/// BD transport stream write sink with embedded FMKV metadata
/// header.
pub struct M2tsStream {
    disc_title: DiscTitle,
    muxer: super::tsmux::TsMuxer<Box<dyn Write + Send>>,
}

impl M2tsStream {
    /// Create for writing PES frames → BD-TS output.
    /// Writes FMKV metadata header, then muxes PES frames into BD transport stream.
    pub fn create(mut writer: impl Write + Send + 'static, title: &DiscTitle) -> io::Result<Self> {
        // Write FMKV metadata header unconditionally. An empty streams
        // array is valid JSON and round-trips fine; skipping the header
        // for a zero-stream title would make the output indistinguishable
        // from a non-FMKV file on read-back (read_header returns
        // Ok(None) → PMT fallback) even though M2tsStream produced it.
        let m = meta::M2tsMeta::from_title(title);
        meta::write_header(&mut writer, &m)?;
        let pids: Vec<u16> = title
            .streams
            .iter()
            .map(|s| match s {
                DiscStream::Video(v) => v.pid,
                DiscStream::Audio(a) => a.pid,
                DiscStream::Subtitle(s) => s.pid,
            })
            .collect();
        let boxed: Box<dyn Write + Send> = Box::new(writer);
        let mut muxer = super::tsmux::TsMuxer::new(boxed, &pids);
        // Declare each video track's codec. This one call decides both the ES
        // framing (HEVC/H.264 arrive length-prefixed and need Annex-B conversion;
        // MPEG-2 and VC-1 are already start-code ES and would be mangled by it)
        // and which parameter-set record parser applies (avcC vs hvcC). Passing
        // the codec rather than a NAL-or-not flag is deliberate: the two facts
        // must never be able to disagree.
        for (i, s) in title.streams.iter().enumerate() {
            if let DiscStream::Video(v) = s {
                muxer.set_video_codec(i, v.codec)?;
            }
        }
        for (i, cp) in title.codec_privates.iter().enumerate() {
            // codec_privates is parallel to streams/pids; ignore any
            // trailing entries that exceed the track count rather than
            // surfacing a track-range error for a benign metadata overrun.
            if i >= pids.len() {
                break;
            }
            if let Some(data) = cp {
                muxer.set_codec_private(i, data.clone())?;
            }
        }
        Ok(Self {
            disc_title: title.clone(),
            muxer,
        })
    }
}

impl crate::pes::Stream for M2tsStream {
    fn read(&mut self) -> io::Result<Option<crate::pes::PesFrame>> {
        // Write-only sink. The m2ts:// read direction is served by
        // `super::resolve::build_m2ts_pipeline` →
        // `PipelinedPesStream`; routing through this type for reads
        // was removed when the highway became the only ingress.
        Err(crate::error::Error::StreamWriteOnly.into())
    }

    fn write(&mut self, frame: &crate::pes::PesFrame) -> io::Result<()> {
        self.muxer
            .write_frame(frame.track, frame.pts, frame.keyframe, &frame.data)
    }

    fn finish(&mut self) -> io::Result<()> {
        self.muxer.finish()
    }

    fn info(&self) -> &crate::disc::DiscTitle {
        &self.disc_title
    }

    fn codec_private(&self, _track: usize) -> Option<Vec<u8>> {
        // Write side doesn't have parsers; codec_private flows in
        // via the title metadata at `create` time and gets baked
        // into the FMKV header. Nothing to surface back here.
        None
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::disc::{
        Codec, ColorSpace, ContentFormat, DiscTitle, FrameRate, HdrFormat, Resolution,
        Stream as DiscStream, VideoStream,
    };
    use crate::pes::{PesFrame, Stream as PesStreamTrait};

    const VIDEO_PID: u16 = 0x1011;

    fn make_title() -> DiscTitle {
        DiscTitle {
            playlist: String::new(),
            playlist_id: 0,
            duration_secs: 0.0,
            size_bytes: 0,
            clips: Vec::new(),
            streams: vec![DiscStream::Video(VideoStream {
                pid: VIDEO_PID,
                codec: Codec::Hevc,
                resolution: Resolution::R1080p,
                frame_rate: FrameRate::F24,
                hdr: HdrFormat::Sdr,
                color_space: ColorSpace::Bt709,
                display_aspect: None,
                secondary: false,
                label: String::new(),
                measured_cicp: None,
            })],
            chapters: Vec::new(),
            extents: Vec::new(),
            content_format: ContentFormat::BdTs,
            codec_privates: vec![Some({
                // Minimal hvcC with one VPS-like array entry.
                let marker: &[u8] = &[0x40, 0x01, 0x0C, 0x01];
                let mut hvcc = vec![0u8; 22];
                hvcc.push(1); // numArrays
                hvcc.push(32);
                hvcc.extend_from_slice(&1u16.to_be_bytes()); // numNalus
                hvcc.extend_from_slice(&(marker.len() as u16).to_be_bytes());
                hvcc.extend_from_slice(marker);
                hvcc
            })],
        }
    }

    fn fake_idr_pes_data() -> Vec<u8> {
        // 4-byte length prefix + NAL: type 19 (IDR_W_RADL).
        let mut nal = vec![(19u8 << 1) & 0x7E, 0x01];
        for i in 0..200 {
            nal.push((i & 0xFF) as u8);
        }
        let mut out = Vec::with_capacity(4 + nal.len());
        out.extend_from_slice(&(nal.len() as u32).to_be_bytes());
        out.extend_from_slice(&nal);
        out
    }

    /// Writer wrapper that shares an Arc<Mutex<Vec<u8>>> so the test can
    /// inspect the bytes after the muxer drops.
    struct SharedSink(std::sync::Arc<std::sync::Mutex<Vec<u8>>>);
    impl Write for SharedSink {
        fn write(&mut self, b: &[u8]) -> io::Result<usize> {
            self.0.lock().unwrap().extend_from_slice(b);
            Ok(b.len())
        }
        fn flush(&mut self) -> io::Result<()> {
            Ok(())
        }
    }

    /// An H.264 track's codec_private is an **avcC** record, not hvcC. The BD-TS
    /// muxer must parse it with the avcC parser and emit the SPS/PPS as Annex-B
    /// parameter sets, or the H.264 elementary stream reaches the player with no
    /// SPS/PPS at all and is undecodable — silently, because frame_count still
    /// advances and the mux reports success.
    ///
    /// Mutation: parse codec_private with hvcc_to_annex_b (the pre-fix behaviour) ->
    /// the parser returns None, no parameter sets are emitted, and this fails.
    #[test]
    fn h264_avcc_parameter_sets_are_emitted_as_annex_b() {
        let sps: &[u8] = &[0x67, 0x42, 0xC0, 0x1E, 0xAB, 0xCD];
        let pps: &[u8] = &[0x68, 0xCE, 0x3C, 0x80];
        // avcC (ISO/IEC 14496-15 §5.3.3.1.2): 5-byte fixed header, then
        // numOfSequenceParameterSets (low 5 bits), each SPS as u16-BE length +
        // bytes, then numOfPictureParameterSets, each PPS likewise.
        let mut avcc = vec![0x01, 0x42, 0xC0, 0x1E, 0xFF];
        avcc.push(0xE0 | 1); // reserved 111b + numSPS = 1
        avcc.extend_from_slice(&(sps.len() as u16).to_be_bytes());
        avcc.extend_from_slice(sps);
        avcc.push(1); // numPPS = 1
        avcc.extend_from_slice(&(pps.len() as u16).to_be_bytes());
        avcc.extend_from_slice(pps);

        let mut title = make_title();
        if let DiscStream::Video(v) = &mut title.streams[0] {
            v.codec = Codec::H264;
        }
        title.codec_privates = vec![Some(avcc)];

        // A length-prefixed IDR NAL, the shape the muxer expects for NAL video.
        let nal: Vec<u8> = vec![0x65, 0x88, 0x84, 0x00, 0x11, 0x22];
        let mut es = (nal.len() as u32).to_be_bytes().to_vec();
        es.extend_from_slice(&nal);

        let shared = std::sync::Arc::new(std::sync::Mutex::new(Vec::<u8>::new()));
        let sink = SharedSink(shared.clone());
        let mut stream = M2tsStream::create(sink, &title).unwrap();
        stream
            .write(&PesFrame {
                coding: None,
                source: None,
                track: 0,
                pts: 0,
                keyframe: true,
                data: es,
                duration_ns: None,
            })
            .unwrap();
        stream.finish().unwrap();
        drop(stream);

        let buf = shared.lock().unwrap().clone();
        assert!(
            buf.windows(sps.len()).any(|w| w == sps),
            "the avcC SPS must reach the transport stream"
        );
        assert!(
            buf.windows(pps.len()).any(|w| w == pps),
            "the avcC PPS must reach the transport stream"
        );
    }

    /// `M2tsStream::create` must opt a VC-1 video track OUT of Annex-B conversion.
    ///
    /// This pins the WIRING in `create`, not just `TsMuxer`'s flag: deleting the
    /// `set_video_codec` loop leaves every TsMuxer-level test passing, because those
    /// drive the muxer directly and set the flag themselves. Only a test that goes
    /// through `create` catches it — and mangling MPEG-2/VC-1 video is silent, since
    /// frame_count still increments and the mux reports success.
    ///
    /// Mutation: remove the `set_video_codec` loop from `create`, or make it declare
    /// Vc1 as a NAL codec -> the ES gains a start code and this fails.
    #[test]
    fn vc1_video_is_wired_to_the_non_nal_path() {
        let mut title = make_title();
        if let DiscStream::Video(v) = &mut title.streams[0] {
            v.codec = Codec::Vc1;
        }
        title.codec_privates = vec![None];

        // Length-prefix SHAPED ES: if the conversion is wrongly applied it rewrites
        // these leading four bytes into a 00 00 00 01 start code.
        let es: Vec<u8> = vec![0x00, 0x00, 0x00, 0x06, 0x0F, 0x12, 0x34, 0x56, 0x78, 0x9A];

        let shared = std::sync::Arc::new(std::sync::Mutex::new(Vec::<u8>::new()));
        let sink = SharedSink(shared.clone());
        let mut stream = M2tsStream::create(sink, &title).unwrap();
        stream
            .write(&PesFrame {
                coding: None,
                source: None,
                track: 0,
                pts: 0,
                keyframe: true,
                data: es.clone(),
                duration_ns: None,
            })
            .unwrap();
        stream.finish().unwrap();
        drop(stream);

        let buf = shared.lock().unwrap().clone();
        assert!(
            buf.windows(es.len()).any(|w| w == &es[..]),
            "VC-1 ES must reach the output verbatim, not converted to Annex-B"
        );
    }

    #[test]
    fn m2ts_stream_forwards_keyframe_to_rai() {
        let title = make_title();
        let shared = std::sync::Arc::new(std::sync::Mutex::new(Vec::<u8>::new()));
        let sink = SharedSink(shared.clone());
        let mut stream = M2tsStream::create(sink, &title).unwrap();
        let frame = PesFrame {
            coding: None,
            source: None,
            track: 0,
            pts: 0,
            keyframe: true,
            data: fake_idr_pes_data(),
            duration_ns: None,
        };
        stream.write(&frame).unwrap();
        stream.finish().unwrap();
        drop(stream);

        let buf = shared.lock().unwrap().clone();

        // Skip FMKV metadata header via meta::read_header.
        let mut cursor = std::io::Cursor::new(&buf);
        let _meta = super::meta::read_header(&mut cursor)
            .unwrap()
            .expect("FMKV header present");
        let header_end = cursor.position() as usize;
        let ts_bytes = &buf[header_end..];

        // Find first PUSI packet on VIDEO_PID; verify RAI in AF flags.
        // chunks_exact drops any partial trailing chunk — only whole
        // 192-byte BD-TS packets are valid, and it avoids OOB indexing on a
        // short final chunk.
        let pkt = ts_bytes
            .chunks_exact(192)
            .find(|p| {
                let h = &p[4..];
                let pid = (((h[1] & 0x1F) as u16) << 8) | h[2] as u16;
                pid == VIDEO_PID && (h[1] & 0x40) != 0
            })
            .expect("video PUSI packet present");
        let h = &pkt[4..];
        let afc = (h[3] >> 4) & 0x03;
        assert!(afc & 0b10 != 0, "AF must be present");
        let af_len = h[4] as usize;
        assert!(af_len >= 1, "AF length must include flags byte");
        let flags = h[5];
        assert_eq!(flags & 0x40, 0x40, "RAI bit set");
    }
}