//! Decode-correctness probe: read an MCAP and, for every
//! `foxglove.CompressedVideo` channel, collect its frames in log_time order,
//! check the first frame is a keyframe (IDR), and decode the entire stream
//! through libavcodec — reporting any decode error. A clean smart-cut output
//! has `first_is_keyframe=true` and decodes every frame with no error.
//!
//! Usage: `cargo run --bin probe -- <file.mcap>`

use std::collections::BTreeMap;

use anyhow::{anyhow, Result};
use cutmerge::{nal, pb};
use ffmpeg_next as ff;
use ff::codec::Id;
use prost::Message as _;

fn main() -> Result<()> {
    ff::init()?;
    ff::log::set_level(ff::log::Level::Error);

    let argv: Vec<String> = std::env::args().collect();
    if argv.get(1).map(|s| s == "records").unwrap_or(false) {
        let path = argv
            .get(2)
            .ok_or_else(|| anyhow!("usage: probe records <file.mcap>"))?;
        return dump_records(path);
    }

    if argv.get(1).map(|s| s == "compare").unwrap_or(false) {
        match (argv.get(2), argv.get(3), argv.get(4), argv.get(5)) {
            (Some(op), Some(ot), Some(ip), Some(it)) => return compare(op, ot, ip, it),
            _ => {
                return Err(anyhow!(
                    "usage: probe compare <out.mcap> <out_topic> <in.mcap> <in_topic>"
                ))
            }
        }
    }

    let path = argv
        .get(1)
        .cloned()
        .ok_or_else(|| anyhow!("usage: probe <file.mcap> | probe records <file.mcap>"))?;
    let data = std::fs::read(&path)?;

    let mut per_topic: BTreeMap<String, Vec<(u64, Vec<u8>)>> = BTreeMap::new();
    for m in mcap::MessageStream::new(&data)? {
        let m = m?;
        let is_video = m
            .channel
            .schema
            .as_ref()
            .map(|s| s.name == "foxglove.CompressedVideo")
            .unwrap_or(false);
        if !is_video {
            continue;
        }
        let cv = pb::foxglove::CompressedVideo::decode(&m.data[..])?;
        per_topic
            .entry(m.channel.topic.clone())
            .or_default()
            .push((m.log_time, cv.data));
    }

    if per_topic.is_empty() {
        println!("no foxglove.CompressedVideo channels found in {path}");
        return Ok(());
    }

    let mut all_ok = true;
    for (topic, mut frames) in per_topic {
        frames.sort_by_key(|(t, _)| *t);
        let first_kf = frames
            .first()
            .map(|(_, d)| nal::is_keyframe(d))
            .unwrap_or(false);
        let (decoded, err) = decode_all(frames.iter().map(|(_, d)| d.as_slice()));
        let ok = first_kf && err.is_none() && decoded == frames.len();
        all_ok &= ok;
        println!(
            "{:<42} frames={:<5} first_is_keyframe={:<5} decoded={:<5} {}",
            topic,
            frames.len(),
            first_kf,
            decoded,
            match &err {
                Some(e) => format!("DECODE ERROR: {e}"),
                None if ok => "OK".to_string(),
                None => "MISMATCH".to_string(),
            }
        );
    }

    if all_ok {
        println!("\nALL VIDEO CHANNELS DECODE CLEANLY");
        Ok(())
    } else {
        Err(anyhow!("one or more video channels failed the decode check"))
    }
}

/// Walk the top-level MCAP records and print a run-length-encoded opcode
/// sequence — the clearest view of the byte-order. mcap-optimiser layout shows
/// all `Chunk`s grouped, then one relocated `MessageIndex` block, then the
/// summary section (Schema, Channel, ChunkIndex, Statistics, SummaryOffset).
fn dump_records(path: &str) -> Result<()> {
    let data = std::fs::read(path)?;
    if data.len() < 9 || &data[..8] != mcap::MAGIC {
        return Err(anyhow!("not an MCAP file (bad magic)"));
    }
    let mut pos = 8usize;
    let mut seq: Vec<(u8, usize)> = Vec::new();
    while pos + 9 <= data.len() {
        let op = data[pos];
        let len = u64::from_le_bytes(data[pos + 1..pos + 9].try_into().unwrap()) as usize;
        match seq.last_mut() {
            Some(last) if last.0 == op => last.1 += 1,
            _ => seq.push((op, 1)),
        }
        pos += 9 + len;
        if op == 0x02 {
            break; // Footer
        }
    }
    println!("{path}:");
    for (op, n) in &seq {
        println!("  {n:>6} x  0x{op:02x} {}", op_name(*op));
    }
    Ok(())
}

fn op_name(op: u8) -> &'static str {
    match op {
        0x01 => "Header",
        0x02 => "Footer",
        0x03 => "Schema",
        0x04 => "Channel",
        0x05 => "Message",
        0x06 => "Chunk",
        0x07 => "MessageIndex",
        0x08 => "ChunkIndex",
        0x09 => "Attachment",
        0x0a => "AttachmentIndex",
        0x0b => "Statistics",
        0x0c => "Metadata",
        0x0d => "MetadataIndex",
        0x0e => "SummaryOffset",
        0x0f => "DataEnd",
        _ => "?",
    }
}

/// Collect a topic's frames as (log_time, optional sensor timestamp, payload).
/// For CompressedVideo the payload is the H.264 Annex-B (so re-encoded frames
/// show up as "payload differs"); otherwise it's the raw message bytes.
fn collect(path: &str, topic: &str) -> Result<Vec<(u64, Option<(i64, i32)>, Vec<u8>)>> {
    let data = std::fs::read(path)?;
    let mut out = Vec::new();
    for m in mcap::MessageStream::new(&data)? {
        let m = m?;
        if m.channel.topic != topic {
            continue;
        }
        let is_video = m
            .channel
            .schema
            .as_ref()
            .map(|s| s.name == "foxglove.CompressedVideo")
            .unwrap_or(false);
        if is_video {
            let cv = pb::foxglove::CompressedVideo::decode(&m.data[..])?;
            out.push((m.log_time, cv.timestamp.map(|t| (t.seconds, t.nanos)), cv.data));
        } else {
            out.push((m.log_time, None, m.data.into_owned()));
        }
    }
    out.sort_by_key(|x| x.0);
    Ok(out)
}

/// Diff an output topic against its source input topic: confirm every output
/// frame's log_time and sensor timestamp match the input, and report how many
/// payloads are byte-identical (passthrough) vs re-encoded.
fn compare(out_path: &str, out_topic: &str, in_path: &str, in_topic: &str) -> Result<()> {
    let out = collect(out_path, out_topic)?;
    let inp = collect(in_path, in_topic)?;
    let by_lt: std::collections::HashMap<u64, usize> =
        inp.iter().enumerate().map(|(i, f)| (f.0, i)).collect();

    let (mut missing, mut ts_mismatch) = (0usize, 0usize);
    let (mut identical, mut differs, mut differs_ts_ok) = (0usize, 0usize, 0usize);
    for f in &out {
        let Some(&j) = by_lt.get(&f.0) else {
            missing += 1;
            continue;
        };
        let g = &inp[j];
        let ts_ok = f.1 == g.1;
        if !ts_ok {
            ts_mismatch += 1;
        }
        if f.2 == g.2 {
            identical += 1;
        } else {
            differs += 1;
            if ts_ok {
                differs_ts_ok += 1;
            }
        }
    }
    let matched = out.len() - missing;
    println!("output {out_topic}  vs  input {in_topic}");
    println!("  output frames:                  {}", out.len());
    println!("  log_time found in input:        {}/{}", matched, out.len());
    println!(
        "  sensor timestamp matches input: {}/{}",
        matched - ts_mismatch,
        matched
    );
    println!("  payload byte-identical:         {identical}  (passthrough)");
    println!(
        "  payload differs:                {differs}  (re-encoded; timestamp still matches: {differs_ts_ok})"
    );
    Ok(())
}

/// Decode every Annex-B frame in order, returning (decoded_picture_count, error).
fn decode_all<'a>(frames: impl Iterator<Item = &'a [u8]>) -> (usize, Option<String>) {
    let dec_codec = match ff::codec::decoder::find(Id::H264) {
        Some(c) => c,
        None => return (0, Some("no H.264 decoder".into())),
    };
    let dec_ctx = ff::codec::context::Context::new_with_codec(dec_codec);
    let mut dec = match dec_ctx.decoder().video() {
        Ok(d) => d,
        Err(e) => return (0, Some(e.to_string())),
    };

    let mut count = 0usize;
    for annexb in frames {
        let mut pkt = ff::codec::packet::Packet::copy(annexb);
        pkt.set_pts(None);
        pkt.set_dts(None);
        if let Err(e) = dec.send_packet(&pkt) {
            return (count, Some(format!("send_packet: {e}")));
        }
        loop {
            let mut f = ff::frame::Video::empty();
            match dec.receive_frame(&mut f) {
                Ok(()) => count += 1,
                Err(ff::Error::Other { errno }) if errno == ff::error::EAGAIN => break,
                Err(ff::Error::Eof) => break,
                Err(e) => return (count, Some(format!("receive_frame: {e}"))),
            }
        }
    }
    let _ = dec.send_eof();
    loop {
        let mut f = ff::frame::Video::empty();
        match dec.receive_frame(&mut f) {
            Ok(()) => count += 1,
            _ => break,
        }
    }
    (count, None)
}
