//! DCVC-class INTER-frame neural video decoder (`wai.neural.dcvc_rt`,
//! `wai.neural.dcvc_fm`).
//!
//! "Native-sink, requires a GPU" was a SPEED claim, not a possibility claim:
//! inter-frame decode is a frame loop that propagates a reference feature, and
//! ONNX Runtime runs that on CPU on any platform. The temporal modeling lives
//! in the sink decoder net (DCVC-RT's implicit temporal context / DCVC-FM's
//! feature modulation); this file orchestrates it:
//!
//!   F ← 0                                  (no temporal context before frame 0)
//!   for each frame f:
//!     y_hat ← hyperprior entropy-decode of frame f      (shared with images)
//!     image, F ← temporal_decoder(y_hat, F)             (F carries frame→frame)
//!     emit image
//!
//! The entropy decode is the exact CompressAI-validated mean-scale hyperprior
//! path (`hyperprior::decode_yhat`); the only video-specific piece is the
//! reference-feature buffer threaded through the loop.
//!
//! Wire format (little-endian):
//!   <IIIIHHBBHH>  H, W, n_frames, fps×1000, Cy, Cz, Sy_log2, Sz_log2, q, Cf
//!   n_frames × <II>   per-frame Lz, Ly
//!   then per frame:   Lz bytes z-bitstream | Ly bytes y-bitstream
//!
//! Sink ONNX: hyper_synthesis.onnx (+ hyper/gauss CDF json) as for the image
//! hyperprior, plus temporal_decoder.onnx:
//!   inputs:  latents      float [1, Cy, Sy, Sy]
//!            ref_feature  float [1, Cf, Sy, Sy]   (zeros for frame 0)
//!   outputs: image        float [1, 3, H, W]
//!            feature      float [1, Cf, Sy, Sy]   (propagated to the next frame)


use ndarray::Array4;
use ort::value::Value;

use super::hyperprior::decode_yhat;
use super::video::DecodedVideo;
use super::{params::ParamSource, DecodeError};

const HEADER: usize = 26;

pub(crate) fn decode(payload: &[u8], params: &ParamSource) -> Result<DecodedVideo, DecodeError> {
    if payload.len() < HEADER {
        return Err(DecodeError::InvalidPayload(format!(
            "dcvc payload {} B < {HEADER} B header", payload.len())));
    }
    let u32a = |o: usize| u32::from_le_bytes(payload[o..o + 4].try_into().unwrap());
    let u16a = |o: usize| u16::from_le_bytes(payload[o..o + 2].try_into().unwrap());
    let (h, w) = (u32a(0) as usize, u32a(4) as usize);
    let n = u32a(8) as usize;
    let fps = u32a(12) as f32 / 1000.0;
    let (cy, cz) = (u16a(16) as usize, u16a(18) as usize);
    let (sy, sz) = (1usize << payload[20], 1usize << payload[21]);
    let q = u16a(22) as u32;
    let cf = u16a(24) as usize;

    let table = HEADER + 8 * n;
    if payload.len() < table {
        return Err(DecodeError::InvalidPayload(format!(
            "dcvc payload {} < header + {n} frame-table entries ({table})", payload.len())));
    }
    let mut lens = Vec::with_capacity(n);
    for i in 0..n {
        let o = HEADER + 8 * i;
        lens.push((u32a(o) as usize, u32a(o + 4) as usize));
    }

    // Load the temporal decoder once; reuse across frames.
    let mut tdec = params.session_sibling("temporal_decoder.onnx")?;

    // Reference feature, propagated frame→frame. Zero before the first frame.
    let mut feat = Array4::<f32>::zeros((1, cf, sy, sy));
    let mut frames_rgb = Vec::with_capacity(n);
    let mut off = table;
    for (i, &(lz, ly)) in lens.iter().enumerate() {
        if off + lz + ly > payload.len() {
            return Err(DecodeError::InvalidPayload(format!(
                "dcvc frame {i} bitstreams out of bounds (need {} at {off})", lz + ly)));
        }
        let z_bits = &payload[off..off + lz];
        let y_bits = &payload[off + lz..off + lz + ly];
        off += lz + ly;

        let y_hat = decode_yhat(z_bits, y_bits, params, q, cz, cy, sz, sy)?;

        let y_v = Value::from_array(y_hat)
            .map_err(|e| DecodeError::Ort(format!("dcvc frame {i} latents: {e}")))?;
        let f_v = Value::from_array(feat.clone())
            .map_err(|e| DecodeError::Ort(format!("dcvc frame {i} ref_feature: {e}")))?;
        let outputs = tdec.run(ort::inputs!["latents" => y_v, "ref_feature" => f_v])
            .map_err(|e| DecodeError::Ort(format!("dcvc frame {i} temporal_decoder: {e}")))?;

        // image → RGB.
        let img = outputs.get("image")
            .ok_or_else(|| DecodeError::Ort("dcvc output 'image' missing".into()))?;
        let iv = img.try_extract_array::<f32>()
            .map_err(|e| DecodeError::Ort(format!("extract image: {e}")))?;
        let iv = iv.view();
        let id = iv.shape();
        if id[2] != h || id[3] != w {
            return Err(DecodeError::Ort(format!(
                "dcvc frame {i} image [1,3,{},{}]; want [1,3,{h},{w}]", id[2], id[3])));
        }
        let data = iv.as_slice().expect("contiguous image");
        let hw = h * w;
        let mut rgb = vec![0u8; hw * 3];
        for p in 0..hw {
            for ch in 0..3 {
                rgb[p * 3 + ch] = (data[ch * hw + p].clamp(0.0, 1.0) * 255.0).round() as u8;
            }
        }
        frames_rgb.push(rgb);

        // feature → next reference (owned copy; the ort output borrows the run).
        let fv = outputs.get("feature")
            .ok_or_else(|| DecodeError::Ort("dcvc output 'feature' missing".into()))?;
        let fview = fv.try_extract_array::<f32>()
            .map_err(|e| DecodeError::Ort(format!("extract feature: {e}")))?;
        let fview = fview.view();
        let fd = fview.shape();
        if fd[1] != cf || fd[2] != sy || fd[3] != sy {
            return Err(DecodeError::Ort(format!(
                "dcvc frame {i} feature [1,{},{},{}]; want [1,{cf},{sy},{sy}]", fd[1], fd[2], fd[3])));
        }
        feat = Array4::from_shape_vec((1, cf, sy, sy), fview.as_slice().expect("contiguous feature").to_vec())
            .map_err(|e| DecodeError::Ort(format!("feature reshape: {e}")))?;
    }

    Ok(DecodedVideo { frames_rgb, width: w as u32, height: h as u32, fps })
}
