use muxide::api::{MuxerBuilder, MuxerConfig, VideoCodec};
use muxide::codec::common::AnnexBNalIter;
use muxide::codec::h264::{annexb_to_avcc, extract_avc_config, is_h264_keyframe};
use muxide::fragmented::{FragmentConfig, FragmentedMuxer};
use std::{
    env,
    fs::{self, File},
    io::{self, Read, Write},
    path::{Path, PathBuf},
};

fn split_access_units_by_aud(annexb: &[u8]) -> Vec<Vec<u8>> {
    let start_code: [u8; 4] = [0, 0, 0, 1];

    let mut aus: Vec<Vec<u8>> = Vec::new();
    let mut current: Vec<u8> = Vec::new();

    for nal in AnnexBNalIter::new(annexb) {
        if nal.is_empty() {
            continue;
        }

        let nal_type = nal[0] & 0x1f;

        // NAL type 9 = Access Unit Delimiter.
        if nal_type == 9 && !current.is_empty() {
            aus.push(current);
            current = Vec::new();
        }

        current.extend_from_slice(&start_code);
        current.extend_from_slice(nal);
    }

    if !current.is_empty() {
        aus.push(current);
    }

    aus
}

fn read_file(path: &Path) -> io::Result<Vec<u8>> {
    let mut f = File::open(path)?;
    let mut buf = Vec::new();
    f.read_to_end(&mut buf)?;
    Ok(buf)
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // Interop check using a *real* H.264 stream generated by ffmpeg.
    // Prereqs: ffmpeg installed and in PATH.
    // Usage:
    //   1) ffmpeg -f lavfi -i testsrc=size=640x480:rate=30 -t 1 \
    //        -pix_fmt yuv420p -c:v libx264 -x264-params aud=1:repeat-headers=1 -f h264 target/interop2/generated.h264
    //   2) cargo run --example interop_from_ffmpeg_h264 -- target/interop2

    let out_dir: PathBuf = env::args_os()
        .nth(1)
        .map(PathBuf::from)
        .unwrap_or_else(|| PathBuf::from("target/interop2"));

    fs::create_dir_all(&out_dir)?;

    let h264_path = out_dir.join("generated.h264");
    let annexb = read_file(&h264_path).map_err(|e| {
        io::Error::new(
            e.kind(),
            format!(
                "failed to read {}: run the ffmpeg command in the example header first",
                h264_path.display()
            ),
        )
    })?;

    let access_units = split_access_units_by_aud(&annexb);
    if access_units.is_empty() {
        return Err(io::Error::new(io::ErrorKind::InvalidData, "no access units found").into());
    }

    let first_key_au = access_units
        .iter()
        .find(|au| is_h264_keyframe(au))
        .ok_or_else(|| {
            io::Error::new(io::ErrorKind::InvalidData, "no keyframe access unit found")
        })?;

    let avc = extract_avc_config(first_key_au).ok_or_else(|| {
        io::Error::new(
            io::ErrorKind::InvalidData,
            "generated stream keyframe missing SPS/PPS",
        )
    })?;

    // 1) muxide non-fragmented MP4.
    let muxide_mp4_path = out_dir.join("muxide.mp4");
    {
        let file = File::create(&muxide_mp4_path)?;
        let config = MuxerConfig::new(640, 480, 30.0);
        let mut muxer = MuxerBuilder::new(file)
            .video(
                VideoCodec::H264,
                config.width,
                config.height,
                config.framerate,
            )
            .build()?;

        for (i, au) in access_units.iter().enumerate() {
            let pts = (i as f64) / 30.0;
            let is_key = is_h264_keyframe(au);
            muxer.write_video(pts, au, is_key)?;
        }

        muxer.finish()?;
    }

    // 2) muxide fragmented MP4, concatenated init+segment so ffprobe can open it.
    let cfg = FragmentConfig {
        width: 640,
        height: 480,
        timescale: 90_000,
        fragment_duration_ms: 1000,
        sps: avc.sps,
        pps: avc.pps,
        vps: None,
        av1_sequence_header: None,
        vp9_config: None,
    };

    let mut fmux = FragmentedMuxer::new(cfg);
    let init = fmux.init_segment();

    let dt = 3000u64;
    for (i, au) in access_units.iter().enumerate() {
        let ticks = (i as u64) * dt;
        let is_key = is_h264_keyframe(au);
        let avcc = annexb_to_avcc(au);
        fmux.write_video(ticks, ticks, &avcc, is_key)?;
    }

    let seg = fmux
        .flush_segment()
        .ok_or_else(|| io::Error::other("fragment did not flush"))?;

    let muxide_fmp4_combined_path = out_dir.join("muxide_fmp4_combined.mp4");
    {
        let mut f = File::create(&muxide_fmp4_combined_path)?;
        f.write_all(&init)?;
        f.write_all(&seg)?;
        f.flush()?;
    }

    eprintln!("Wrote: {}", muxide_mp4_path.display());
    eprintln!("Wrote: {}", muxide_fmp4_combined_path.display());

    Ok(())
}
