//! Decoder for `wai.neural.bmshj2018` (v0.4 wire format).
//!
//! Wire format:
//!   <IIIIBB>  H, W, L, quality, C, S_log2
//!   L bytes   CompressAI rANS bitstream
//!
//! Decode pipeline:
//!   1. Load `cdfs_q<quality>.json` from the same dir as decoder.onnx
//!      (cached per-process).
//!   2. Range-decode L bytes → C*S*S int32 symbols (channel-major order).
//!   3. Add per-channel medians → float32 [1, C, S, S] latents.
//!   4. Run g_s through ONNX → float32 [1, 3, H, W] in [0, 1].

use std::collections::VecDeque;
use std::sync::{Arc, Mutex, OnceLock};

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

use super::{rans::RansDecoder, params::ParamSource, DecodeError};

/// Decoded image: RGB row-major u8, H × W × 3.
#[derive(Debug, Clone)]
pub struct DecodedImage {
    pub rgb: Vec<u8>,
    pub width: u32,
    pub height: u32,
}

#[derive(Debug, Clone, Deserialize)]
struct CdfBlob {
    channels:   usize,
    cdf_length: Vec<usize>,
    offset:     Vec<i32>,
    medians:    Vec<f32>,
    cdf:        Vec<Vec<i32>>,
}

/// How many parsed CDF tables the process keeps. A sink decodes against a few
/// parameter sets at a time; the bound keeps a long-running sink that meets many
/// distinct tables (updated sets, unpinned directories that change) from growing
/// without limit.
const CDF_CACHE_CAP: usize = 8;

/// Parsed CDF tables by the SHA-256 of the bytes they were parsed from, oldest
/// first; at most [`CDF_CACHE_CAP`] entries.
type CdfCache = VecDeque<(String, Arc<CdfBlob>)>;

fn cdf_cache() -> &'static Mutex<CdfCache> {
    static CACHE: OnceLock<Mutex<CdfCache>> = OnceLock::new();
    CACHE.get_or_init(|| Mutex::new(VecDeque::new()))
}

/// The entry for `key`, moved to the newest end, or `None`.
fn cache_get(cache: &mut CdfCache, key: &str) -> Option<Arc<CdfBlob>> {
    let at = cache.iter().position(|(k, _)| k == key)?;
    let hit = cache.remove(at)?;
    let blob = Arc::clone(&hit.1);
    cache.push_back(hit);
    Some(blob)
}

/// Insert `blob` under `key`, evicting the least recently used entries beyond
/// [`CDF_CACHE_CAP`]. A decode still holding an evicted table keeps its `Arc`.
fn cache_put(cache: &mut CdfCache, key: String, blob: Arc<CdfBlob>) {
    if cache_get(cache, &key).is_some() {
        return;
    }
    cache.push_back((key, blob));
    while cache.len() > CDF_CACHE_CAP {
        cache.pop_front();
    }
}

/// The CDF tables for `quality`, parsed once per distinct content: the cache is
/// keyed by the SHA-256 of the bytes read, so a changed file, or another set's
/// tables, is never answered from a stale entry.
fn load_cdfs_for(params: &ParamSource, quality: u32) -> Result<Arc<CdfBlob>, DecodeError> {
    let name = format!("cdfs_q{quality}.json");
    let raw = params.sibling(&name)?;
    let key = crate::pin::sha256_hex(&raw);
    if let Some(b) = cache_get(&mut cdf_cache().lock().unwrap(), &key) {
        return Ok(b);
    }
    let blob: Arc<CdfBlob> = Arc::new(serde_json::from_slice(&raw)
        .map_err(|e| DecodeError::InvalidPayload(format!("cdf parse {name}: {e}")))?);
    cache_put(&mut cdf_cache().lock().unwrap(), key, Arc::clone(&blob));
    Ok(blob)
}

pub(crate) fn decode(payload: &[u8], params: &ParamSource) -> Result<DecodedImage, DecodeError> {
    if payload.len() < 18 {
        return Err(DecodeError::InvalidPayload(
            format!("payload too short ({} B); need ≥ 18 B header", payload.len())));
    }
    let h = u32::from_le_bytes(payload[0..4].try_into().unwrap()) as usize;
    let w = u32::from_le_bytes(payload[4..8].try_into().unwrap()) as usize;
    let l = u32::from_le_bytes(payload[8..12].try_into().unwrap()) as usize;
    let q = u32::from_le_bytes(payload[12..16].try_into().unwrap());
    let c = payload[16] as usize;
    let s = 1usize << payload[17];
    if payload.len() != 18 + l {
        return Err(DecodeError::InvalidPayload(format!(
            "payload length {} != 18 + L={l}", payload.len())));
    }
    let rans_bytes = &payload[18..18 + l];

    let cdfs = load_cdfs_for(params, q)?;
    if cdfs.channels != c {
        return Err(DecodeError::InvalidPayload(format!(
            "envelope C={c} != cdf channels={}", cdfs.channels)));
    }

    // Range-decode C * S * S symbols, channel-major (matches
    // EntropyBottleneck._build_indexes flatten order).
    let mut latents = Array4::<f32>::zeros((1, c, s, s));
    let lat_slice = latents.as_slice_mut().unwrap();
    let mut dec = RansDecoder::new(rans_bytes)
        .map_err(|e| DecodeError::InvalidPayload(e.into()))?;
    let mut idx = 0;
    for ch in 0..c {
        let cdf     = &cdfs.cdf[ch];
        let cdf_len = cdfs.cdf_length[ch];
        let offset  = cdfs.offset[ch];
        let median  = cdfs.medians[ch];
        for _ in 0..(s * s) {
            let sym = dec.decode(cdf, cdf_len, offset);
            lat_slice[idx] = sym as f32 + median;
            idx += 1;
        }
    }

    let mut sess = params.session_model("decoder.onnx")?;
    let input_name  = sess.inputs().first().map(|i| i.name().to_string())
        .unwrap_or_else(|| "latents".into());
    let output_name = sess.outputs().first().map(|o| o.name().to_string())
        .unwrap_or_else(|| "image".into());
    let latents_v = Value::from_array(latents)
        .map_err(|e| DecodeError::Ort(format!("latents tensor: {e}")))?;
    let outputs = sess.run(ort::inputs![input_name => latents_v])
        .map_err(|e| DecodeError::Ort(format!("session.run: {e}")))?;
    let out = outputs.get(output_name.as_str()).ok_or_else(||
        DecodeError::Ort(format!("output {output_name} missing")))?;
    let view = out.try_extract_array::<f32>()
        .map_err(|e| DecodeError::Ort(format!("extract image: {e}")))?;
    let view = view.view();
    let dims = view.shape();
    if dims.len() != 4 || dims[1] != 3 || dims[2] != h || dims[3] != w {
        return Err(DecodeError::Ort(format!(
            "unexpected output shape {dims:?}; want [1, 3, {h}, {w}]")));
    }
    let data = view.as_slice().expect("contiguous");
    let mut rgb = vec![0u8; h * w * 3];
    for y in 0..h {
        for x in 0..w {
            let pi = y * w + x;
            for ch in 0..3 {
                let v = data[ch * h * w + pi].clamp(0.0, 1.0);
                rgb[pi * 3 + ch] = (v * 255.0).round() as u8;
            }
        }
    }
    Ok(DecodedImage { rgb, width: w as u32, height: h as u32 })
}

#[cfg(test)]
mod tests {
    use super::*;

    fn blob(channels: usize) -> Arc<CdfBlob> {
        Arc::new(CdfBlob { channels, cdf_length: vec![], offset: vec![], medians: vec![], cdf: vec![] })
    }

    /// The CDF cache holds at most CDF_CACHE_CAP tables, evicting the least
    /// recently used; a table a decode still holds outlives its eviction.
    #[test]
    fn cdf_cache_is_bounded() {
        let mut c = CdfCache::new();
        let first = blob(0);
        cache_put(&mut c, "k0".into(), Arc::clone(&first));
        for i in 1..3 * CDF_CACHE_CAP {
            cache_put(&mut c, format!("k{i}"), blob(i));
            // k1 is read on every round, so it stays.
            assert!(i < 2 || cache_get(&mut c, "k1").is_some());
        }
        assert_eq!(c.len(), CDF_CACHE_CAP);
        assert!(cache_get(&mut c, "k0").is_none(), "the oldest unused table is evicted");
        assert_eq!(cache_get(&mut c, "k1").unwrap().channels, 1);
        assert_eq!(first.channels, 0, "an evicted table a decode holds is still valid");
        assert_eq!(Arc::strong_count(&first), 1);
    }
}
