//! Smart-cut re-encoder for the leading partial GOP at a cut boundary.
//!
//! When a segment's first video frame is not a keyframe, the frames from the
//! window start up to the next keyframe reference a keyframe that lies *before*
//! the window — so on their own they are undecodable. This helper decodes them
//! (primed by the frames from that prior keyframe) and re-encodes just the
//! in-window ("emit") frames into a fresh, self-contained H.264 GOP whose first
//! frame is an IDR. The caller passes every complete GOP after the next
//! keyframe through untouched, so only this short run is ever re-encoded.
//!
//! B-frames are not used (matching slimcap / `foxglove.CompressedVideo`), so
//! decode order == input order == display order and frames map 1:1 by position.

use anyhow::{anyhow, bail, Context, Result};
use ffmpeg_next as ff;
use ff::codec::Id;
use ff::format::Pixel;
use ff::software::scaling::{Context as Scaler, Flags};
use ff::util::picture;
use ff::util::rational::Rational;

/// One input frame to the leading-GOP re-encoder, in stream order.
pub struct FrameIn<'a> {
    pub annexb: &'a [u8],
    /// false => priming (decode only, not emitted); true => re-encode and emit.
    pub emit: bool,
}

/// Decode `frames` in order (priming first) and re-encode the `emit` frames
/// into a fresh self-contained H.264 GOP (IDR forced on the first emit frame).
/// Returns the new Annex-B bytestream for each emit frame, in input order.
pub fn reencode_lead_gop(frames: &[FrameIn], crf: u32, preset: &str) -> Result<Vec<Vec<u8>>> {
    let emit_count = frames.iter().filter(|f| f.emit).count();
    if emit_count == 0 {
        return Ok(Vec::new());
    }

    // 1. Decode every frame to YUV, in order (priming frames seed references).
    let dec_codec =
        ff::codec::decoder::find(Id::H264).ok_or_else(|| anyhow!("no H.264 decoder"))?;
    let dec_ctx = ff::codec::context::Context::new_with_codec(dec_codec);
    let mut decoder = dec_ctx.decoder().video().context("open H.264 decoder")?;

    let mut decoded: Vec<ff::frame::Video> = Vec::with_capacity(frames.len());
    for f in frames {
        let mut pkt = ff::codec::packet::Packet::copy(f.annexb);
        pkt.set_pts(None);
        pkt.set_dts(None);
        decoder.send_packet(&pkt).context("decoder send_packet")?;
        drain_decoder(&mut decoder, &mut decoded)?;
    }
    decoder.send_eof().context("decoder send_eof")?;
    drain_decoder(&mut decoder, &mut decoded)?;

    if decoded.len() != frames.len() {
        bail!(
            "H.264 decode produced {} frames for {} inputs (stream gap?)",
            decoded.len(),
            frames.len()
        );
    }
    // Emit frames are the trailing `emit_count` (priming frames precede them).
    let emit_yuv = &decoded[frames.len() - emit_count..];

    // 2. Re-encode the emit frames into one fresh GOP.
    let (w, h) = (emit_yuv[0].width(), emit_yuv[0].height());
    if w % 2 != 0 || h % 2 != 0 {
        bail!("yuv420p requires even dimensions, got {w}x{h}");
    }

    let enc_codec =
        ff::codec::encoder::find_by_name("libx264").ok_or_else(|| anyhow!("no libx264 encoder"))?;
    let enc_ctx = ff::codec::context::Context::new_with_codec(enc_codec);
    let mut enc = enc_ctx.encoder().video()?;
    enc.set_width(w);
    enc.set_height(h);
    enc.set_format(Pixel::YUV420P);
    enc.set_time_base(Rational(1, 1_000_000));
    enc.set_max_b_frames(0);
    enc.set_frame_rate(Some(Rational(30, 1)));

    let mut opts = ff::Dictionary::new();
    opts.set("preset", preset);
    opts.set("crf", &crf.to_string());
    // Disable the encoder's own keyframe cadence (we force the head IDR) and
    // keep each IDR self-contained with inline SPS/PPS, Annex-B framed.
    opts.set(
        "x264-params",
        "repeat-headers=1:annexb=1:keyint=9999:min-keyint=1:scenecut=0",
    );
    let mut enc = enc.open_with(opts).context("open libx264 encoder")?;

    let mut out: Vec<Vec<u8>> = Vec::with_capacity(emit_count);
    let mut scaler: Option<(Pixel, Scaler)> = None;
    for (i, src) in emit_yuv.iter().enumerate() {
        let src_fmt = src.format();
        let rebuild = scaler.as_ref().map(|(f, _)| *f != src_fmt).unwrap_or(true);
        if rebuild {
            let s = Scaler::get(src_fmt, w, h, Pixel::YUV420P, w, h, Flags::BILINEAR)
                .context("build swscale context")?;
            scaler = Some((src_fmt, s));
        }
        let mut yuv = ff::frame::Video::new(Pixel::YUV420P, w, h);
        scaler
            .as_mut()
            .unwrap()
            .1
            .run(src, &mut yuv)
            .context("swscale to yuv420p")?;
        yuv.set_pts(Some(i as i64));
        if i == 0 {
            yuv.set_kind(picture::Type::I);
        }
        enc.send_frame(&yuv).context("encoder send_frame")?;
        // Drain ready packets after every frame: libx264 buffers internally and
        // send_frame returns EAGAIN if we don't keep it drained. bframes=0 keeps
        // packets in input order, 1:1 with emit frames.
        drain_encoder(&mut enc, &mut out)?;
    }
    enc.send_eof().context("encoder send_eof")?;
    drain_encoder(&mut enc, &mut out)?;

    if out.len() != emit_count {
        bail!(
            "re-encode produced {} packets for {} emit frames",
            out.len(),
            emit_count
        );
    }
    Ok(out)
}

fn drain_decoder(
    decoder: &mut ff::codec::decoder::Video,
    out: &mut Vec<ff::frame::Video>,
) -> Result<()> {
    loop {
        let mut frame = ff::frame::Video::empty();
        match decoder.receive_frame(&mut frame) {
            Ok(()) => out.push(frame),
            Err(ff::Error::Other { errno }) if errno == ff::error::EAGAIN => break,
            Err(ff::Error::Eof) => break,
            Err(e) => return Err(e).context("decoder receive_frame"),
        }
    }
    Ok(())
}

fn drain_encoder(enc: &mut ff::codec::encoder::Video, out: &mut Vec<Vec<u8>>) -> Result<()> {
    loop {
        let mut pkt = ff::codec::packet::Packet::empty();
        match enc.receive_packet(&mut pkt) {
            Ok(()) => out.push(pkt.data().unwrap_or_default().to_vec()),
            Err(ff::Error::Other { errno }) if errno == ff::error::EAGAIN => break,
            Err(ff::Error::Eof) => break,
            Err(e) => return Err(e).context("encoder receive_packet"),
        }
    }
    Ok(())
}
