//! Self-sufficient JPEG AI (ISO/IEC 6048) decode: a real codestream → RGB, with
//! NO per-image schedule fixtures. Every per-image quantity is derived from the
//! bitstream — the residual coding mask (`sigma_skip > 382`, jpeg-ai skip_mode),
//! the σ surfaces, the bucket indexes, the rvs additions; everything else is a
//! **model constant** the sink supplies in `model_dir` (the σ-net weights, the
//! factorized-bottleneck CDFs, the per-channel dequant gain, the normative
//! entropy tables, the LSBS/route consts, and the ONNX nets).
//!
//! The entropy + integer σ path runs through the integer operator set, by the same
//! lowering and driver as the exact decoder (`crate::jpegai_ops`,
//! integer-operators §13.6.1), so both decode paths derive the same symbols; the learned transforms
//! (hyper-decoder, MCM value nets, synthesis, eICCI post-filters) are
//! sink-supplied ONNX run via `ort`. The reference dispatch uses this path for
//! `wai.image.jpegai` when the sink has no integer weight bundle for the exact
//! decoder (`jpegai_decode`); the e2e test in `tests/jpegai_e2e.rs` drives it against the
//! reference decoder's PNG.
//!
//! Sizing note: the bundled ONNX nets are traced for 256×256 (`synthesis16`,
//! the 4×4 hyper input). This decoder targets that traced model; other sizes
//! need nets traced for them. `img_width/height` from the header are asserted.
//!
//! Scope: the bundle is one configuration — tools_0 (`model_id` 0), the
//! high-operation-point synthesis, one eICCI model pair, and a dequant scaler
//! precomputed for `beta_displacement_log` = 0. `decode` refuses any codestream
//! that [`crate::jpegai_decode::check_codestream`] refuses (the exact decoder's
//! header conditions), so both `wai.image.jpegai` paths accept the same codestreams.
//! This path runs in float through ONNX Runtime and is not decode-equivalent.

use std::path::Path;

use ort::session::Session;
use ort::value::Value;

use crate::neural::image::DecodedImage;
use crate::neural::{params::ParamSource, DecodeError};
use crate::neural::jpegai_bitstream::{parse_ae_substream, parse_codestream};
use crate::neural::jpegai_filters::{efe_linear_plane, efe_nonlinear_plane, lef_apply, EfeLinearPlane};
use crate::neural::jpegai_header::parse_headers;
use crate::jpegai_ops::{self, Plane, ScalePath};

/// JPEG AI decode error.
#[derive(Debug)]
#[non_exhaustive]
pub enum JpegAiError {
    Model(String),
    /// A pinned decode needed a file its verified parameter set does not contain
    /// (`set_incomplete`, SPEC §3.1).
    Pin(crate::container::PinRefusal),
    Bitstream(String),
    Unsupported(String),
}
impl std::fmt::Display for JpegAiError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Model(e) => write!(f, "jpeg-ai model: {e}"),
            Self::Pin(r) => write!(f, "jpeg-ai model: {}: {r}", r.code()),
            Self::Bitstream(e) => write!(f, "jpeg-ai bitstream: {e}"),
            Self::Unsupported(e) => write!(f, "jpeg-ai unsupported: {e}"),
        }
    }
}
impl std::error::Error for JpegAiError {}

type R<T> = Result<T, JpegAiError>;
fn m<E: std::fmt::Display>(e: E) -> JpegAiError { JpegAiError::Model(e.to_string()) }

// ---------------------------------------------------------------------------
// Model bundle — the sink-supplied constants, loaded once from `model_dir`.
// ---------------------------------------------------------------------------

/// The model constants a JPEG AI decode needs, loaded from a model directory.
/// All are fixed per checkpoint (image-invariant); none is per-image.
pub struct JpegAiModel {
    source: ParamSource,
    /// The scale path and the entropy decodes' tables, through the integer operator set
    /// (integer-operators §13.6.1): from `hyper_scale_decoder_int_weights.json`,
    /// `jpegai_normative_tables.json`, `me-tANS/bounds.csv`,
    /// `jpegai_factorized_cdfs.json` and the gains of `model_constants.json`.
    scale: ScalePath,
    consts: serde_json::Value,                // model_constants.json (gain_mlog, rvs, lef)
    scaler_vec: [Vec<f32>; 2],                 // GainUnit per-channel divisor [y, uv]
    rvs_dq: [Vec<i64>; 2],                     // residual dequant tables, grfs-gated [flag0, flag1]
    lsbs_t0: Vec<i64>,
    lsbs_t1: Vec<i64>,
}

fn param_err(e: DecodeError) -> JpegAiError {
    match e {
        DecodeError::PriorMismatch(r) => JpegAiError::Pin(r),
        other => m(other),
    }
}
fn jload(src: &ParamSource, name: &str) -> R<serde_json::Value> {
    serde_json::from_slice(&src.sibling(name).map_err(param_err)?).map_err(|e| m(format!("{name}: {e}")))
}
fn i64s(v: &serde_json::Value) -> Vec<i64> {
    v.as_array().map(|a| a.iter().filter_map(|n| n.as_i64()).collect()).unwrap_or_default()
}
fn f32s(v: &serde_json::Value) -> Vec<f32> {
    v.as_array().map(|a| a.iter().filter_map(|n| n.as_f64().map(|x| x as f32)).collect()).unwrap_or_default()
}

impl JpegAiModel {
    /// Load every model constant from a model directory (the sink's installed
    /// JPEG AI model bundle).
    pub fn load(dir: impl AsRef<Path>) -> R<Self> {
        Self::load_from(ParamSource::Unpinned(dir.as_ref().to_path_buf()))
    }

    /// Load every model constant from a parameter source: a model directory
    /// (unpinned), or the parameter set a pin verified, from which every file —
    /// the ONNX networks included — is read as the bytes that were hashed. A file
    /// the verified set does not contain is [`JpegAiError::Pin`] (`set_incomplete`).
    pub fn load_from(source: ParamSource) -> R<Self> {
        let src = &source;
        let sigma_weights = jload(src, "hyper_scale_decoder_int_weights.json")?;
        let consts = jload(src, "model_constants.json")?;
        let norm = jload(src, "jpegai_normative_tables.json")?;
        let norm_pmf: Vec<Vec<i64>> =
            norm["pmf"].as_array().ok_or_else(|| m("normative pmf"))?.iter().map(i64s).collect();
        let bounds_raw = src.sibling("me-tANS/bounds.csv").map_err(param_err)?;
        let bounds: Vec<i64> = std::str::from_utf8(&bounds_raw)
            .map_err(|e| m(format!("bounds.csv: {e}")))?
            .lines().filter(|l| !l.trim().is_empty())
            .map(|l| l.trim().parse().map_err(|e| m(format!("bounds: {e}")))).collect::<R<_>>()?;

        let fac = jload(src, "jpegai_factorized_cdfs.json")?;
        let plane_cdfs = |p: &str| -> Vec<i64> {
            fac[p].as_array().map(|a| a.iter().flat_map(i64s).collect()).unwrap_or_default()
        };
        let fac_per_channel = fac["per_channel"].as_i64().ok_or_else(|| m("per_channel"))?;

        let sv = jload(src, "jpegai_scaler_vec.json")?;
        let scaler_vec = [f32s(&sv["y"]), f32s(&sv["uv"])];
        let rvs_dq = [i64s(&consts["rvs_dq_t_flag0"]), i64s(&consts["rvs_dq_t_flag1"])];

        let dq = jload(src, "jpegai_dequant_consts.json")?;
        let z_offsets = [i64s(&dq["z_offsets"]["y"]), i64s(&dq["z_offsets"]["uv"])];
        let lsbs_t0 = i64s(&dq["lsbs_t0"]);
        let lsbs_t1 = i64s(&dq["lsbs_t1"]);

        // The scale path and the entropy decodes go through the integer operator set, the
        // same lowering and driver as the exact decoder's (integer-operators §13.6.1). This
        // bundle writes the same members another way; one that is not §13.6.1's (a
        // hyper-latent offset other than 31, a load check failed) is refused.
        if z_offsets.iter().any(|o| o.iter().any(|&v| v != jpegai_ops::Z_OFFSET)) {
            return Err(m("model_mismatch: a hyper-latent offset is not 31 (integer-operators §13.6.1)"));
        }
        let mismatch = |r: wai_ops::Refusal| m(format!("{}: the scale path fails a load check of integer-operators §13.6.1, or does not meet acc32", r.code.as_str()));
        let net = wai_ops::lower::jpegai::JpegaiNet {
            norm: wai_ops::lower::jpegai::Norm {
                pmf: norm_pmf,
                bounds,
                per_channel: fac_per_channel,
                cdfs_y: plane_cdfs("y"),
                cdfs_uv: plane_cdfs("uv"),
                gain_y: i64s(&consts["tools_0"]["gain_mlog_y"]),
                gain_uv: i64s(&consts["tools_0"]["gain_mlog_uv"]),
            },
            hs_y: jpegai_ops::scale_net_of_blocks(&sigma_weights["y"]).map_err(mismatch)?,
            hs_uv: jpegai_ops::scale_net_of_blocks(&sigma_weights["uv"]).map_err(mismatch)?,
        };
        let scale = jpegai_ops::load_net(net).map_err(mismatch)?;

        Ok(Self { source, scale, consts, scaler_vec, rvs_dq, lsbs_t0, lsbs_t1 })
    }

    fn sess(&self, rel: &str) -> R<Session> {
        self.source.session_sibling(rel).map_err(param_err)
    }
}

// ---------------------------------------------------------------------------
// ONNX helpers (mirror the reference trace).
// ---------------------------------------------------------------------------

fn run1(s: &mut Session, x: Vec<f32>, c: usize, h: usize, w: usize) -> R<(Vec<f32>, Vec<usize>)> {
    let name = s.inputs().first().map(|i| i.name().to_string()).ok_or_else(|| m("no input"))?;
    let arr = ndarray::Array4::from_shape_vec((1, c, h, w), x).map_err(m)?;
    let v = Value::from_array(arr).map_err(m)?;
    let out = s.run(ort::inputs![name.as_str() => v]).map_err(m)?;
    let view = out[0].try_extract_array::<f32>().map_err(m)?;
    Ok((view.iter().copied().collect(), view.shape().to_vec()))
}

fn run_mcm(s: &mut Session, p: Vec<f32>, pc: usize, prev: Vec<f32>, vc: usize) -> R<Vec<f32>> {
    let names: Vec<String> = s.inputs().iter().map(|i| i.name().to_string()).collect();
    let pa = ndarray::Array4::from_shape_vec((1, pc, 8, 8), p).map_err(m)?;
    let mut feeds: Vec<(std::borrow::Cow<'static, str>, ort::session::SessionInputValue<'static>)> =
        vec![(names[0].clone().into(), Value::from_array(pa).map_err(m)?.into())];
    if names.len() > 1 {
        let va = ndarray::Array4::from_shape_vec((1, vc, 8, 8), prev).map_err(m)?;
        feeds.push((names[1].clone().into(), Value::from_array(va).map_err(m)?.into()));
    }
    let out = s.run(feeds).map_err(m)?;
    Ok(out[0].try_extract_array::<f32>().map_err(m)?.iter().copied().collect())
}

/// [C,16,16] → 4 stage slices [C,8,8] in decode order (TL, BR, TR, BL).
fn down_shuffle(x: &[f32], c: usize) -> [Vec<f32>; 4] {
    let phase = |py: usize, px: usize| -> Vec<f32> {
        let mut out = Vec::with_capacity(c * 64);
        for ch in 0..c { for y in 0..8 { for xx in 0..8 {
            out.push(x[ch * 256 + (2 * y + py) * 16 + 2 * xx + px]);
        }}}
        out
    };
    [phase(0, 0), phase(1, 1), phase(0, 1), phase(1, 0)]
}

/// Inverse: stage slices (TL,BR,TR,BL) → [C,16,16].
fn up_shuffle(st: &[Vec<f32>; 4], c: usize) -> Vec<f32> {
    let mut out = vec![0f32; c * 256];
    for (si, (py, px)) in [(0usize, (0usize, 0usize)), (1, (1, 1)), (2, (0, 1)), (3, (1, 0))] {
        for ch in 0..c { for y in 0..8 { for xx in 0..8 {
            out[ch * 256 + (2 * y + py) * 16 + 2 * xx + px] = st[si][ch * 64 + y * 8 + xx];
        }}}
    }
    out
}

/// LSBS scaling: σ_skip-driven table lookup.
fn lsbs(yhat: &[f32], resi_dq: &[f32], sigma_skip: &[i64], c: usize, t0: &[i64], t1: &[i64]) -> Vec<f32> {
    let mut out = vec![0f32; c * 256];
    for ch in 0..c {
        let mut blocks = [[0i64; 2]; 2];
        for y in 0..16 { for x in 0..16 { blocks[y / 8][x / 8] += sigma_skip[ch * 256 + y * 16 + x]; } }
        for y in 0..16 { for x in 0..16 {
            let lk = ((blocks[y / 8][x / 8] + 32) >> 6).clamp(0, 4351) as usize;
            let i = ch * 256 + y * 16 + x;
            let mean_full = yhat[i] - resi_dq[i];
            out[i] = yhat[i] + (t0[lk] as f32 * mean_full + t1[lk] as f32 * resi_dq[i]) / 8192.0;
        }}
    }
    out
}

// ---------------------------------------------------------------------------
// The decode.
// ---------------------------------------------------------------------------

/// Decode a JPEG AI codestream to RGB against a model bundle. Every per-image
/// value is derived from `codestream`; `model` supplies only checkpoint
/// constants. Targets the 256×256 traced model (see module note).
pub fn decode(codestream: &[u8], model: &JpegAiModel) -> R<DecodedImage> {
    let subs = parse_codestream(codestream)
        .ok_or_else(|| JpegAiError::Bitstream("parse codestream".into()))?;
    let payload = |mk: u16| -> R<&[u8]> {
        subs.iter().find(|s| s.marker == mk).map(|s| s.payload)
            .ok_or_else(|| JpegAiError::Bitstream(format!("missing marker {mk:#06x}")))
    };

    // This path hard-codes the same configuration as the exact decoder: the tools_0
    // tables and nets (traced at 256×256), the β = 0 dequant scaler
    // (`jpegai_scaler_vec.json`), the eICCI model pair below, the high-operation-point
    // synthesis. A codestream that asks for anything else is refused here rather than
    // decoded to the wrong image.
    crate::jpegai_decode::check_codestream(codestream)
        .map_err(|why| JpegAiError::Unsupported(why.into()))?;
    let hdr = parse_headers(payload(0xFF82)?, payload(0xFF83)?)
        .map_err(|e| JpegAiError::Bitstream(format!("headers: {e:?}")))?;
    let (h, w) = (hdr.img_width as usize, hdr.img_height as usize);
    if (h, w) != (256, 256) {
        return Err(JpegAiError::Unsupported(format!("{h}x{w}; bundled nets are 256×256-traced")));
    }

    // --- 1–3. SOZ → z → the scale graph → σ, mask → the residuals: the scale path and the
    // entropy decodes through the integer operator set (integer-operators §13.6.1). ---
    let (mem_z, ends_z) = parse_ae_substream(payload(0xFF88)?, 1)
        .ok_or_else(|| JpegAiError::Bitstream("SOZ substream".into()))?;
    let mut sigma_skip: Vec<Vec<i64>> = Vec::new();
    let mut sigma_final: Vec<Vec<i64>> = Vec::new();
    let mut z_net_planes: Vec<Vec<i64>> = Vec::new();
    let mut lk_planes: Vec<Vec<usize>> = Vec::new(); // per-block σ-skip "likely" (σ rvs + dequant)
    let mut resi: Vec<Vec<i16>> = Vec::new();
    for (pi, (plane, marker)) in [(Plane::Luma, 0xFF89u16), (Plane::Chroma, 0xFF8A)].into_iter().enumerate() {
        let (mem, ends) = parse_ae_substream(payload(marker)?, 1)
            .ok_or_else(|| JpegAiError::Bitstream("residual substream".into()))?;
        let grfs: Vec<bool> = hdr.models[pi].grfs_channel_flag.iter().map(|&f| f == 1).collect();
        let ints = jpegai_ops::plane(&model.scale, plane, &mem_z[..ends_z[0] as usize], &mem[..ends[0] as usize], &grfs)
            .map_err(|r| JpegAiError::Bitstream(format!("{}: an entropy substream ends before the decode it carries", r.code.as_str())))?;
        sigma_final.push(ints.v.iter().zip(&ints.rvs).map(|(v, r)| v + r).collect());
        lk_planes.push(ints.likely.iter().map(|&l| l as usize).collect());
        // Every symbol is a 16-bit value (E8's wrap16).
        resi.push(ints.y_q.iter().map(|&q| q as i16).collect());
        z_net_planes.push(ints.z);
        sigma_skip.push(ints.v);
    }

    // --- 4. dequant: resi × rvs_dq[grfs[ch]][lk(block)] / 2^16 ÷ scaler_vec[ch].
    // rvs_dq + lk are derived; scaler_vec is the only model parameter (GainUnit). ---
    let mut resi_dq: Vec<Vec<f32>> = Vec::new();
    for pi in 0..2 {
        let sv = &model.scaler_vec[pi];
        let grfs = &hdr.models[pi].grfs_channel_flag;
        let plane: Vec<f32> = resi[pi].iter().enumerate().map(|(i, &r)| {
            let ch = i / 256;
            let dq = &model.rvs_dq[grfs[ch] as usize];
            r as f32 * dq[lk_planes[pi][i]] as f32 / 65536.0 / sv[ch]
        }).collect();
        resi_dq.push(plane);
    }

    // --- 5. hyper decoder (ONNX) → ψ ---
    let mut psi: Vec<Vec<f32>> = Vec::new();
    for (pi, (plane, c)) in [("y", 160usize), ("uv", 96)].iter().enumerate() {
        let mut s = model.sess(&format!("model/CCS_SGMM/tools_0/model_{plane}/common_modules/hyper_decoder.onnx"))?;
        let zf: Vec<f32> = z_net_planes[pi].iter().map(|&v| v as f32).collect();
        psi.push(run1(&mut s, zf, *c, 4, 4)?.0);
    }

    // --- 6. value path: Y via MCM stages, UV via ψ up-shuffle ---
    let resi_y_st = down_shuffle(&resi_dq[0], 160);
    let mut yhat_st: Vec<Vec<f32>> = Vec::new();
    for s_i in 0..4 {
        let mut st = model.sess(&format!("model/CCS_SGMM/tools_0/model_y/common_modules/MCM/stage{s_i}.onnx"))?;
        let p: Vec<f32> = psi[0][s_i * 160 * 64..(s_i + 1) * 160 * 64].to_vec();
        let prev_c = 160 * yhat_st.len().max(1);
        let prev = if yhat_st.is_empty() { vec![0f32; 160 * 64] } else { yhat_st.iter().flatten().copied().collect() };
        let mu = run_mcm(&mut st, p, 160, prev, prev_c)?;
        yhat_st.push(resi_y_st[s_i].iter().zip(&mu).map(|(a, b)| a + b).collect());
    }
    let yhat_y = up_shuffle(&[yhat_st[0].clone(), yhat_st[1].clone(), yhat_st[2].clone(), yhat_st[3].clone()], 160);
    let chf = |i: usize| psi[1][i * 96 * 64..(i + 1) * 96 * 64].to_vec();
    let mu_uv = up_shuffle(&[chf(0), chf(1), chf(2), chf(3)], 96);
    let yhat_uv: Vec<f32> = resi_dq[1].iter().zip(&mu_uv).map(|(a, b)| a + b).collect();

    // --- 7. LSBS ---
    let y_ls = lsbs(&yhat_y, &resi_dq[0], &sigma_skip[0], 160, &model.lsbs_t0, &model.lsbs_t1);
    let uv_ls = lsbs(&yhat_uv, &resi_dq[1], &sigma_skip[1], 96, &model.lsbs_t0, &model.lsbs_t1);

    // --- 8. synthesis (ONNX, 16×16-traced) ---
    let mut sy = model.sess("model16/CCS_SGMM/tools_0/model_y/synthesis16.onnx")?;
    let (mut yp, _) = run1(&mut sy, y_ls.clone(), 160, 16, 16)?;
    let mut suv = model.sess("model16/CCS_SGMM/tools_0/model_uv/synthesis16.onnx")?;
    let mut uv_in = y_ls;
    uv_in.extend_from_slice(&uv_ls);
    let (uvpix, _) = run1(&mut suv, uv_in, 256, 16, 16)?;
    let mut up: Vec<f32> = uvpix[0..h * w].to_vec();
    let mut vp: Vec<f32> = uvpix[h * w..2 * h * w].to_vec();

    // --- 9. EFElinear (parameters from the tool header) ---
    let deint = |q: u64| (q as f32 - 32767.0) / 2048.0;
    for (ci, wq, w2q, mean_q) in [
        (1usize, hdr.efe_linear.weights_u[0], hdr.efe_linear.weights_u2[0], hdr.efe_linear.mean_u_q100),
        (2, hdr.efe_linear.weights_v[0], hdr.efe_linear.weights_v2[0], hdr.efe_linear.mean_v_q100),
    ] {
        if !hdr.efe_linear_enabled { break; }
        let mean = mean_q as f32 / 100.0;
        let p = EfeLinearPlane { mean, w_uv: [deint(wq) + 1.0; 4], bias: [mean; 4], w_y: [deint(w2q); 4] };
        let src = if ci == 1 { &up } else { &vp };
        let out = efe_linear_plane(&yp, src, h, w, &p);
        if ci == 1 { up = out } else { vp = out }
    }

    // --- 10. eICCI (ONNX wrappers; [0,1] domain) ---
    if hdr.icci_enabled {
        let mut yuv01 = Vec::with_capacity(3 * h * w);
        for pl in [&yp, &up, &vp] { yuv01.extend(pl.iter().map(|&v| v / 255.0)); }
        let crop = |data: &[f32], shape: &[usize], ch: usize| -> Vec<f32> {
            let (oh, ow) = (shape[2], shape[3]);
            let mut out = Vec::with_capacity(h * w);
            for y in 0..h { for x in 0..w { out.push(data[(ch * oh + y) * ow + x] * 255.0); } }
            out
        };
        let mut si_y = model.sess("e2e/pfchain/icci_y_m6.onnx")?;
        let (fy, fysh) = run1(&mut si_y, yuv01.clone(), 3, h, w)?;
        let mut si_uv = model.sess("e2e/pfchain/icci_uv_m3.onnx")?;
        let (fuv, fuvsh) = run1(&mut si_uv, yuv01, 3, h, w)?;
        yp = crop(&fy, &fysh, 0);
        up = crop(&fuv, &fuvsh, 0);
        vp = crop(&fuv, &fuvsh, 1);
    }

    // --- 11. EFEnonlinear (parameters from the tool header) ---
    if hdr.efe_nonlinear_enabled && hdr.efn_v_enabled {
        let wq: Vec<i64> = hdr.efn_weights.iter().map(|&v| v as i64).collect();
        vp = efe_nonlinear_plane(&yp, &vp,
            &[wq[0], wq[1], wq[2], wq[3], wq[4], wq[5], wq[6], wq[7]],
            hdr.efn_min_luma_q100 as i64, hdr.efn_max_luma_q100 as i64);
    }

    // --- 12. LEF (op-point tables at header model_id; σ channel from header) ---
    if hdr.lef_enabled {
        let mid = hdr.model_id as usize;
        let thr_t = i64s(&model.consts["lef_thr"][mid]);
        let mag_t = f32s(&model.consts["lef_mag"][mid]);
        let c0 = hdr.lef_ch_idx as usize * 256;
        yp = lef_apply(&yp, &sigma_final[0][c0..c0 + 256], h, w,
            [thr_t[0], thr_t[1], thr_t[2]], [mag_t[0], mag_t[1], mag_t[2]]);
    }

    // --- 13. BT.709 → RGB ---
    let (kr, kg, kb) = (0.2126f32, 0.7152f32, 0.0722f32);
    let (kby, kry) = (1.8556f32, 1.5748f32);
    let mut rgb = vec![0u8; h * w * 3];
    for i in 0..h * w {
        let y01 = yp[i] / 255.0;
        let u01 = up[i] / 255.0 - 0.5;
        let v01 = vp[i] / 255.0 - 0.5;
        let px = [
            y01 + kry * v01,
            y01 - (kb * kby / kg) * u01 - (kr * kry / kg) * v01,
            y01 + kby * u01,
        ];
        for (c, &val) in px.iter().enumerate() {
            rgb[i * 3 + c] = (val * 255.0).round_ties_even().clamp(0.0, 255.0) as u8;
        }
    }
    Ok(DecodedImage { rgb, width: w as u32, height: h as u32 })
}
