//! Byte-exact **mANS** (masked-ANS) entropy decoder for JPEG-AI — a faithful port
//! of the ISO reference `cpp_exts/mans/decompressor.cpp`. Two interleaved ANS states,
//! MSB-first reads from a rolling 64-bit window (`read`/`flush`), sentinel-bit init,
//! and a group-of-4 inbound→flush→outbound schedule. Pure integer ⇒ byte-identical
//! on every machine. Backs the entropy stage of `wai.image.jpegai`.
//!
//! `decode_sgm` (main y-latent): per position a distribution `sigma_idx` selects one
//! of the 32 precomputed transition tables; `mask=false` ⇒ the position is skipped
//! (value 0) and does not consume the stream. The state threads across calls, so a
//! single [`MansDecoder`] decodes a whole stream over many `decode_sgm` invocations.

// On the JPEG AI exact decode path: no platform math-library call (see clippy.toml).
#![deny(clippy::disallowed_methods)]

/// Low 56..63-bit masks for `flush()` (BIT_MASK_HIGH in the reference).
const BIT_MASK_HIGH: [u64; 8] = [
    0x00ff_ffff_ffff_ffff, 0x01ff_ffff_ffff_ffff, 0x03ff_ffff_ffff_ffff, 0x07ff_ffff_ffff_ffff,
    0x0fff_ffff_ffff_ffff, 0x1fff_ffff_ffff_ffff, 0x3fff_ffff_ffff_ffff, 0x7fff_ffff_ffff_ffff,
];

/// The backward-consumed 64-bit bit window + the two interleaved ANS states.
#[derive(Clone, Copy)]
struct BitStream {
    ptr_end: usize, // index into `memory` (one past the current window's low byte block)
    head: u64,
    bit_pos: u32,
    state1: u8,
    state2: u8,
}

impl BitStream {
    fn le64(mem: &[u8], at: usize) -> u64 {
        let mut b = [0u8; 8];
        b.copy_from_slice(&mem[at..at + 8]);
        u64::from_le_bytes(b)
    }
    /// Read `step` bits (MSB-first) out of the window.
    fn read(&mut self, step: u32) -> u64 {
        self.bit_pos -= step;
        let result = self.head >> self.bit_pos;
        self.head ^= result << self.bit_pos;
        result
    }
    /// Refill the window from earlier bytes (moving backward).
    fn flush(&mut self, mem: &[u8]) {
        self.ptr_end -= (7 ^ (self.bit_pos >> 3)) as usize;
        self.head = Self::le64(mem, self.ptr_end) & BIT_MASK_HIGH[(self.bit_pos & 7) as usize];
        self.bit_pos |= 56;
    }
    fn init_backward(&mut self, mem: &[u8], ptr: usize) {
        self.ptr_end = ptr - 8;
        self.head = Self::le64(mem, self.ptr_end);
        self.bit_pos = 63 ^ self.head.leading_zeros();
        self.head ^= 1u64 << self.bit_pos;
    }
    fn init_with_states(&mut self, mem: &[u8], ptr: usize) {
        self.init_backward(mem, ptr);
        self.state1 = self.read(8) as u8;
        self.state2 = self.read(8) as u8;
        self.flush(mem);
    }
}

const MAX_Z: usize = 63;

/// A byte-exact mANS decoder over one memory buffer (single-thread streams).
pub struct MansDecoder {
    memory: Vec<u8>,          // 8 zero bytes prepended, then the stream bytes
    stream: BitStream,        // persists across decode calls (state threads)
    trans_y: Vec<u32>,        // num_distributions * 256 flat transition table (sgm)
    bounds_y: Vec<u8>,        // per-distribution outbound bound (sgm)
    z_pre: [u16; 512],        // transitionsZPreprocess (built once, for the fac path)
    trans_z: [u32; 256],      // per-row Z transition table (rebuilt per channel)
}

impl MansDecoder {
    /// `mem` = the raw stream bytes; `offsets` = per-thread end offsets
    /// (`cumsum(thread_sizes)`). Only single-thread (`offsets.len()==1`) is supported.
    pub fn new(mem: &[u8], offsets: &[u32]) -> Self {
        assert_eq!(offsets.len(), 1, "only single-thread streams supported");
        let cap = mem.len();
        let mut memory = vec![0u8; cap + 8];
        memory[8..8 + cap].copy_from_slice(mem);
        let thread_off = offsets[0] as usize + 8;
        let mut stream = BitStream { ptr_end: 0, head: 0, bit_pos: 0, state1: 0, state2: 0 };
        stream.init_with_states(&memory, thread_off);
        // transitionsZPreprocess: for i in 511..1, nBits grows so (i<<nBits) reaches [256,512)
        let mut z_pre = [0u16; 512];
        let mut nbits: u32 = 0;
        for i in (1..512usize).rev() {
            if (i << nbits) < 256 { nbits += 1; }
            z_pre[i] = (((i << nbits) ^ 256) as u16) | ((nbits as u16) << 8);
        }
        MansDecoder { memory, stream, trans_y: Vec::new(), bounds_y: Vec::new(), z_pre, trans_z: [0u32; 256] }
    }

    /// Install the 32 sgm transition tables (`num_dist × 256` u32) + per-dist bounds.
    pub fn set_sgm_transitions(&mut self, transitions: &[u32], bounds: &[u8]) {
        self.trans_y = transitions.to_vec();
        self.bounds_y = bounds.to_vec();
    }

    #[inline]
    fn t_y(&self, dist: usize, state: u8) -> u32 {
        self.trans_y[dist * 256 + state as usize]
    }

    #[inline]
    fn decode_y_inbound(&mut self, which: u8, sigma: &[u8], out: &mut [i16], index: usize) {
        let transition = self.t_y(sigma[index] as usize, if which == 1 { self.stream.state1 } else { self.stream.state2 });
        let ns = ((transition >> 16) as u64 | self.stream.read(transition >> 24)) as u8;
        if which == 1 { self.stream.state1 = ns; } else { self.stream.state2 = ns; }
        out[index] = transition as i16; // low 16 bits, signed = the symbol
    }

    #[inline]
    fn decode_y_outbound(&mut self, sigma: &[u8], out: &mut [i16], index: usize) {
        let bound = self.bounds_y[sigma[index] as usize] as i32;
        if out[index] as i32 + bound == 0 {
            let step = if self.stream.read(1) != 0 { 3 } else { 16 };
            let value_temp = self.stream.read(step) as i32;
            let sign = value_temp & 1;
            out[index] = (((value_temp >> 1) + (bound - sign)) ^ (-sign)) as i16;
            self.stream.flush(&self.memory);
        }
    }

    /// Decode one span: `out[i] = symbol` where `mask[i]`, else left untouched (0).
    /// State threads across calls. Faithful port of `decodeSingleThread`.
    pub fn decode_sgm(&mut self, sigma: &[u8], out: &mut [i16], mask: &[bool]) {
        let length = sigma.len();
        let mut index = 0;
        // main groups of 4: state1/state2/state1/state2, flush between in/out bound
        while index + 3 < length {
            if mask[index] { self.decode_y_inbound(1, sigma, out, index); }
            if mask[index + 1] { self.decode_y_inbound(2, sigma, out, index + 1); }
            if mask[index + 2] { self.decode_y_inbound(1, sigma, out, index + 2); }
            if mask[index + 3] { self.decode_y_inbound(2, sigma, out, index + 3); }
            self.stream.flush(&self.memory);
            if mask[index] { self.decode_y_outbound(sigma, out, index); }
            if mask[index + 1] { self.decode_y_outbound(sigma, out, index + 1); }
            if mask[index + 2] { self.decode_y_outbound(sigma, out, index + 2); }
            if mask[index + 3] { self.decode_y_outbound(sigma, out, index + 3); }
            index += 4;
        }
        if length & 2 != 0 {
            if mask[index] {
                self.decode_y_inbound(1, sigma, out, index);
                self.decode_y_outbound(sigma, out, index);
            }
            if mask[index + 1] {
                self.decode_y_inbound(2, sigma, out, index + 1);
                self.decode_y_outbound(sigma, out, index + 1);
            }
            index += 2;
            self.stream.flush(&self.memory);
        }
        if length & 1 != 0 {
            if mask[index] {
                self.decode_y_inbound(1, sigma, out, index);
                self.decode_y_outbound(sigma, out, index);
            }
            self.stream.flush(&self.memory);
        }
    }

    // ---- fac (factorized z hyper-latent) path ----

    /// Build the per-row Z transition table from a channel's cumulative `cdfs`
    /// (`MAX_Z` entries). Faithful port of `setZDistributions`.
    fn set_z_distributions(&mut self, cdfs: &[u8]) {
        let mut curr = 0usize;
        for i in 0..MAX_Z {
            let c = cdfs[i] as usize;
            let mut k = c - curr; // = count for symbol i
            for j in curr..c {
                self.trans_z[j] = ((self.z_pre[k] as u32) << 16) | i as u32;
                k += 1;
            }
            curr = c;
        }
        self.trans_z[255] = ((self.z_pre[1] as u32) << 16) | MAX_Z as u32;
    }

    #[inline]
    fn decode_z_inbound(&mut self, which: u8, out: &mut [u8], index: usize) {
        let transition = self.trans_z[if which == 1 { self.stream.state1 } else { self.stream.state2 } as usize];
        let ns = ((transition >> 16) as u64 | self.stream.read(transition >> 24)) as u8;
        if which == 1 { self.stream.state1 = ns; } else { self.stream.state2 = ns; }
        out[index] = transition as u8; // low 8 bits = symbol (0..=MAX_Z)
    }

    #[inline]
    fn decode_z_outbound(&mut self, out: &mut [u8], index: usize) {
        if out[index] as usize == MAX_Z {
            out[index] = self.stream.read(6) as u8;
        }
    }

    fn decode_row(&mut self, cdfs: &[u8], out: &mut [u8]) {
        self.set_z_distributions(cdfs);
        let length = out.len();
        let mut index = 0;
        while index + 3 < length {
            self.decode_z_inbound(1, out, index);
            self.decode_z_inbound(2, out, index + 1);
            self.decode_z_inbound(1, out, index + 2);
            self.decode_z_inbound(2, out, index + 3);
            self.decode_z_outbound(out, index);
            self.decode_z_outbound(out, index + 1);
            self.decode_z_outbound(out, index + 2);
            self.decode_z_outbound(out, index + 3);
            self.stream.flush(&self.memory);
            index += 4;
        }
        if length & 2 != 0 {
            self.decode_z_inbound(1, out, index);
            self.decode_z_outbound(out, index);
            self.decode_z_inbound(2, out, index + 1);
            self.decode_z_outbound(out, index + 1);
            index += 2;
        }
        if length & 1 != 0 {
            self.decode_z_inbound(1, out, index);
            self.decode_z_outbound(out, index);
        }
        self.stream.flush(&self.memory);
    }

    /// Decode the factorized z hyper-latent: `channels` rows, `size_per_channel`
    /// symbols each, each row using its own `MAX_Z`-entry cumulative CDF (`cdfs` is
    /// `channels × MAX_Z`). Faithful port of `decode_factorize` (single-thread).
    pub fn decode_factorize(&mut self, cdfs: &[u8], out: &mut [u8], channels: usize, size_per_channel: usize) {
        for i in 0..channels {
            let row = &cdfs[i * MAX_Z..(i + 1) * MAX_Z];
            let (a, b) = (i * size_per_channel, (i + 1) * size_per_channel);
            let mut buf = vec![0u8; size_per_channel];
            self.decode_row(row, &mut buf);
            out[a..b].copy_from_slice(&buf);
        }
    }
}

/// The reference's symbol ordering for a distribution of `len` bins: alternating
/// outward from the centre (`mid, mid-1, mid+1, mid-2, …`), which is what pairs each
/// CDF slot with its signed residual value.
fn sequence(len: usize) -> Vec<i64> {
    let mid = len / 2;
    let mut out = vec![0i64; len];
    let (mut even, mut odd) = (mid as i64, mid as i64 - 1);
    for (i, v) in out.iter_mut().enumerate() {
        if i % 2 == 0 { *v = even; even += 1; } else { *v = odd; odd -= 1; }
    }
    out
}

/// Build the 256-entry `sgm` decode table per σ-distribution from the normative
/// integer `(pmf, cdf)` rows — the integer-only port of `utils.get_decode_transitions`.
///
/// Each entry packs `(nBits << 24) | (next_base << 16) | (symbol & 0xffff)`, the layout
/// [`MansDecoder::decode_sgm`] consumes. Each CDF row is split into a first and second
/// half (`c - (c >> 1)` and `(c >> 1) + 128`) so the two interleaved ANS states draw
/// from disjoint slot ranges.
///
/// Having this on the integer path means the entropy stage needs only the normative
/// PMF tables — no precomputed transition fixture, and no dependency on the
/// float/ONNX-side module.
pub fn build_decode_transitions(pmf: &[Vec<i64>], cdf: &[Vec<i64>]) -> Vec<[u32; 256]> {
    const MASS_BITS: u32 = 8;
    let mass = 1i64 << MASS_BITS;
    let width = pmf[0].len();
    let seq: Vec<i64> = sequence(width).iter().map(|&v| v - (width as i64) / 2).collect();
    // get_clz(512): the number of doublings that bring `state` into [256, 512).
    let clz = |state: i64| -> i64 {
        let (mut bits, mut s) = (0i64, state.max(1));
        while s < 256 { s <<= 1; bits += 1; }
        bits
    };
    (0..pmf.len())
        .map(|i| {
            let mut symbols = [0i64; 256];
            let mut state_next = [0i64; 256];
            let cdf_first: Vec<i64> = cdf[i].iter().map(|&c| c - (c >> 1)).collect();
            let cdf_second: Vec<i64> = cdf[i].iter().map(|&c| (c >> 1) + 128).collect();
            let (mut col, mut cdf_curr) = (0usize, 0i64);
            while cdf_curr < mass {
                let p = pmf[i][col];
                let d_first = cdf_first[col + 1] - cdf_first[col];
                for (k, slot) in (cdf_first[col]..cdf_first[col + 1]).enumerate() {
                    symbols[slot as usize] = seq[col];
                    state_next[slot as usize] = p + k as i64;
                }
                for (k, slot) in (cdf_second[col]..cdf_second[col + 1]).enumerate() {
                    symbols[slot as usize] = seq[col];
                    state_next[slot as usize] = p + d_first + k as i64;
                }
                cdf_curr = cdf[i][col + 1];
                col += 1;
            }
            let mut table = [0u32; 256];
            for s in 0..256 {
                let bits = clz(state_next[s]);
                let renormed = (state_next[s] << bits) ^ mass;
                table[s] = ((bits as u32) << 24)
                    | ((renormed as u32) << 16)
                    | ((symbols[s] as u16) as u32);
            }
            table
        })
        .collect()
}

/// Cumulative rows for [`build_decode_transitions`] from the normative PMF tables.
pub fn cdf_from_pmf(pmf: &[Vec<i64>]) -> Vec<Vec<i64>> {
    pmf.iter()
        .map(|row| {
            let mut c = vec![0i64; row.len() + 1];
            for (i, &p) in row.iter().enumerate() { c[i + 1] = c[i] + p; }
            c
        })
        .collect()
}
