//! The NNC form (`extensions/prior-carriage.md` §12): an ISO/IEC 15938-17 bitstream of NNR
//! units, read and written for the subset §12.1 states — a start unit, one model parameter
//! set, at most one topology unit, and one compressed data unit per tensor, each of payload
//! type `NNR_PT_INT` (integers coded as §12.4 codes them, [`super::cabac`]) or
//! `NNR_PT_RAW_FLOAT` (32-bit floats, raw). Everything outside the subset is refused with
//! `nnc_not_read`, naming what it met.

use super::cabac;
use super::{refuse, DType, Refusal, Tensor, TensorSet};
use crate::strict_json::{self, quote, Json};

/// NNR unit types.
pub const NNR_STR: u32 = 0;
pub const NNR_MPS: u32 = 1;
pub const NNR_LPS: u32 = 2;
pub const NNR_TPL: u32 = 3;
pub const NNR_QNT: u32 = 4;
pub const NNR_NDU: u32 = 5;
pub const NNR_AGG: u32 = 6;
/// Compressed data unit payload types.
pub const NNR_PT_INT: u32 = 0;
pub const NNR_PT_FLOAT: u32 = 1;
pub const NNR_PT_RAW_FLOAT: u32 = 2;
pub const NNR_PT_BLOCK: u32 = 3;
/// The topology storage format "unrecognised", whose topology is a string; the writer puts
/// the metadata there (§12.3).
pub const NNR_TPL_UNREC: u32 = 0;
/// The topology storage formats this reader reads, each carrying its topology as a string.
pub const TPL_STRING_FORMATS: [u32; 3] = [0, 3, 4];
/// The unary length the writer codes with: `cabac_unary_length_minus1`.
pub const UNARY_LENGTH_MINUS1: u32 = 10;
/// The most elements one tensor may have at this reader.
pub const MAX_ELEMENTS: u64 = 1 << 28;

struct Bits<'a> {
    b: &'a [u8],
    bit: usize,
}

impl Bits<'_> {
    fn u(&mut self, n: u32) -> Result<u64, Refusal> {
        let mut v = 0u64;
        for _ in 0..n {
            let byte = *self.b.get(self.bit / 8).ok_or_else(|| Refusal::new("nnc_truncated", "a unit ends inside its header"))?;
            v = (v << 1) | u64::from((byte >> (7 - self.bit % 8)) & 1);
            self.bit += 1;
        }
        Ok(v)
    }
    fn ue(&mut self, mut k: u32) -> Result<u64, Refusal> {
        let mut x: u64 = 0;
        loop {
            if k >= 40 {
                return refuse("nnc_malformed", "an exponential-Golomb code longer than 40 bits");
            }
            if self.u(1)? == 1 {
                break;
            }
            x += 1u64 << k;
            k += 1;
        }
        Ok(x + self.u(k)?)
    }
    fn align(&mut self) -> Result<(), Refusal> {
        if self.u(1)? != 1 {
            return refuse("nnc_malformed", "a byte alignment that does not begin with a 1 bit");
        }
        while !self.bit.is_multiple_of(8) {
            if self.u(1)? != 0 {
                return refuse("nnc_malformed", "a byte alignment padded with a 1 bit");
            }
        }
        Ok(())
    }
    fn st(&mut self) -> Result<String, Refusal> {
        if !self.bit.is_multiple_of(8) {
            return refuse("nnc_malformed", "a string that does not begin on a byte");
        }
        let at = self.bit / 8;
        let Some(n) = self.b[at.min(self.b.len())..].iter().position(|&c| c == 0) else {
            return refuse("nnc_truncated", "a string with no terminating zero byte");
        };
        let s = String::from_utf8(self.b[at..at + n].to_vec()).map_err(|_| Refusal::new("nnc_malformed", "a string that is not UTF-8"))?;
        self.bit += (n + 1) * 8;
        Ok(s)
    }
}

#[derive(Default)]
struct W {
    b: Vec<u8>,
    bit: usize,
}

impl W {
    fn u(&mut self, n: u32, v: u64) {
        for i in (0..n).rev() {
            if self.bit.is_multiple_of(8) {
                self.b.push(0);
            }
            if (v >> i) & 1 == 1 {
                let last = self.b.len() - 1;
                self.b[last] |= 1 << (7 - self.bit % 8);
            }
            self.bit += 1;
        }
    }
    fn ue(&mut self, mut k: u32, mut x: u64) {
        while x >= (1u64 << k) {
            self.u(1, 0);
            x -= 1u64 << k;
            k += 1;
        }
        self.u(1, 1);
        self.u(k, x);
    }
    fn align(&mut self) {
        self.u(1, 1);
        while !self.bit.is_multiple_of(8) {
            self.u(1, 0);
        }
    }
    fn st(&mut self, s: &str) {
        debug_assert!(self.bit.is_multiple_of(8));
        self.b.extend_from_slice(s.as_bytes());
        self.b.push(0);
        self.bit = self.b.len() * 8;
    }
}

/// Prefix `body` (a unit's header and payload, after its size) with `nnr_unit_size`: the
/// 15-bit form when the whole unit is under 32768 bytes, else the 31-bit form (§12.2).
fn sized(body: Vec<u8>) -> Result<Vec<u8>, Refusal> {
    let mut w = W::default();
    if body.len() + 2 < 32768 {
        w.u(1, 0);
        w.u(15, (body.len() + 2) as u64);
    } else if body.len() + 4 < (1usize << 31) {
        w.u(1, 1);
        w.u(31, (body.len() + 4) as u64);
    } else {
        return refuse("value_out_of_range", "an NNR unit past 2^31 − 1 bytes");
    }
    w.b.extend_from_slice(&body);
    Ok(w.b)
}

/// The common start of a unit header: its type and two flags, both zero but
/// `independently_decodable_flag`.
fn unit_head(w: &mut W, ty: u32) {
    w.u(6, u64::from(ty));
    w.u(1, 1);
    w.u(1, 0);
}

/// The canonical metadata text a topology unit carries (§12.3): a JSON object of the
/// metadata's strings, keys in ascending byte order.
pub fn metadata_text(set: &TensorSet) -> String {
    let m: Vec<String> = set.metadata.iter().map(|(k, v)| format!("{}:{}", quote(k), quote(v))).collect();
    format!("{{{}}}", m.join(","))
}

/// The canonical bitstream of `set` (§11.3, §12.3).
pub fn write(set: &TensorSet) -> Result<Vec<u8>, Refusal> {
    let set = set.clone().sorted();
    let mut out = Vec::new();
    // NNR_STR: general_profile_idc 0.
    let mut w = W::default();
    unit_head(&mut w, NNR_STR);
    w.u(8, 0);
    out.extend(sized(w.b)?);
    // NNR_MPS: topology carried when there is metadata; nothing else.
    let tpl = !set.metadata.is_empty();
    let mut w = W::default();
    unit_head(&mut w, NNR_MPS);
    w.u(1, u64::from(tpl));
    w.u(4, 0); // sparsification, pruning, unification, decomposition performance map
    w.u(3, 0); // quantization method flags
    w.u(1, 0); // topology indexed reference
    w.u(7, 0); // reserved
    w.align();
    out.extend(sized(w.b)?);
    if tpl {
        let text = metadata_text(&set);
        if text.contains('\0') {
            return refuse("header_malformed", "metadata that holds a zero byte");
        }
        let mut w = W::default();
        unit_head(&mut w, NNR_TPL);
        w.u(8, u64::from(NNR_TPL_UNREC));
        w.u(8, 0); // topology_compression_format: raw
        w.st(&text);
        out.extend(sized(w.b)?);
    }
    for (i, t) in set.tensors.iter().enumerate() {
        if i > 0 && set.tensors[i - 1].name == t.name {
            return refuse("tensor_malformed", format!("{:?} named twice", t.name));
        }
        if t.name.is_empty() || t.name.contains('\0') || t.shape.is_empty() || t.shape.contains(&0) {
            return refuse("tensor_malformed", format!("{:?} cannot be an NNC tensor (a name without a zero byte, at least one dimension, none 0)", t.name));
        }
        let n = t.count().filter(|&n| n <= MAX_ELEMENTS).ok_or_else(|| Refusal::new("value_out_of_range", format!("{:?}: more than 2^28 elements", t.name)))?;
        let pt = match t.dtype {
            d if d.is_profile_integer() => NNR_PT_INT,
            DType::F32 => NNR_PT_RAW_FLOAT,
            d => return refuse("dtype_unknown", format!("{:?}: {d:?} is carried by neither NNR_PT_INT nor NNR_PT_RAW_FLOAT", t.name)),
        };
        if t.data.len() as u64 != t.dtype.bytes_for(n).unwrap_or(u64::MAX) {
            return refuse("offset_mismatch", format!("{:?}: {} bytes for its shape and dtype", t.name, t.data.len()));
        }
        let mut w = W::default();
        unit_head(&mut w, NNR_NDU);
        w.u(5, u64::from(pt));
        w.u(1, 0); // multiple topology elements
        w.u(1, 0); // decompressed data format present
        w.u(1, 1); // input parameters present
        w.st(&t.name);
        if pt == NNR_PT_INT {
            w.u(1, 0); // dq_flag
        }
        w.u(1, 1); // tensor_dimensions_flag
        w.u(1, 1); // cabac_unary_length_flag
        w.u(4, 0); // compressed_parameter_types
        w.ue(1, t.shape.len() as u64);
        for &d in &t.shape {
            w.ue(7, d);
        }
        w.u(8, u64::from(UNARY_LENGTH_MINUS1));
        if t.shape.len() > 1 {
            w.u(4, 0); // scan_order
        }
        w.align();
        let mut body = w.b;
        if pt == NNR_PT_INT {
            let v = t.ints().expect("an integer tensor");
            body.extend(cabac::encode_ints(&v, UNARY_LENGTH_MINUS1 + 1)?);
        } else {
            body.extend_from_slice(&t.data);
        }
        out.extend(sized(body)?);
    }
    Ok(out)
}

/// Read an NNC bitstream (§12.2): the subset of §12.1, every unit's size held to what it
/// carries. Tensors of `NNR_PT_INT` come back as `I32`, of `NNR_PT_RAW_FLOAT` as `F32`; a
/// topology unit of the unrecognised format whose text is a JSON object of strings gives
/// the metadata.
pub fn read(b: &[u8]) -> Result<TensorSet, Refusal> {
    let mut set = TensorSet::default();
    let mut at = 0usize;
    let (mut seen_mps, mut seen_tpl, mut seen_ndu, mut first) = (false, false, false, true);
    while at < b.len() {
        let mut h = Bits { b: &b[at..], bit: 0 };
        let long = h.u(1)?;
        let size = h.u(if long == 1 { 31 } else { 15 })? as usize;
        if size < 3 || at + size > b.len() {
            return refuse("nnc_truncated", format!("a unit of {size} bytes at byte {at} of {}", b.len()));
        }
        let unit = &b[at..at + size];
        let mut h = Bits { b: unit, bit: h.bit };
        let ty = h.u(6)? as u32;
        let _independently_decodable = h.u(1)?;
        let partial = h.u(1)?;
        if first && ty != NNR_STR {
            return refuse("nnc_malformed", "the first unit is not NNR_STR");
        }
        if partial == 1 {
            return refuse("nnc_not_read", format!("unit type {ty} with a partial data counter"));
        }
        let header_done = |h: &Bits<'_>, what: &str| -> Result<(), Refusal> {
            if h.bit != unit.len() * 8 {
                return refuse("nnc_malformed", format!("{what} carries {} bits and its size says {}", h.bit, unit.len() * 8));
            }
            Ok(())
        };
        match ty {
            NNR_STR => {
                if !first {
                    return refuse("nnc_malformed", "a second NNR_STR");
                }
                h.u(8)?; // general_profile_idc
                header_done(&h, "NNR_STR")?;
            }
            NNR_MPS => {
                if seen_mps || seen_ndu {
                    return refuse("nnc_malformed", "NNR_MPS twice, or after a data unit");
                }
                seen_mps = true;
                let _topology = h.u(1)?;
                let maps = h.u(4)?;
                let quant = h.u(3)?;
                let indexed = h.u(1)?;
                if h.u(7)? != 0 {
                    return refuse("nnc_malformed", "reserved bits of NNR_MPS are not zero");
                }
                if maps != 0 {
                    return refuse("nnc_not_read", "a performance map in NNR_MPS");
                }
                if indexed != 0 {
                    return refuse("nnc_not_read", "topology elements referenced by index");
                }
                if quant != 0 {
                    h.u(3)?; // mps_qp_density
                    h.u(13)?; // mps_quantization_parameter
                }
                h.align()?;
                header_done(&h, "NNR_MPS")?;
            }
            NNR_TPL => {
                if seen_tpl || !seen_mps {
                    return refuse("nnc_malformed", "NNR_TPL twice, or before NNR_MPS");
                }
                seen_tpl = true;
                let fmt = h.u(8)? as u32;
                let comp = h.u(8)?;
                if !TPL_STRING_FORMATS.contains(&fmt) || comp != 0 {
                    return refuse("nnc_not_read", format!("topology storage format {fmt}, compression {comp}"));
                }
                let text = h.st()?;
                header_done(&h, "NNR_TPL")?;
                if fmt == NNR_TPL_UNREC
                    && let Ok(Json::Obj(m)) = strict_json::read(text.as_bytes())
                    && m.iter().all(|(_, v)| v.as_str().is_some())
                {
                    for (k, v) in m {
                        set.metadata.insert(k, v.as_str().expect("checked").to_string());
                    }
                }
            }
            NNR_NDU => {
                if !seen_mps {
                    return refuse("nnc_malformed", "a data unit before NNR_MPS");
                }
                seen_ndu = true;
                set.tensors.push(ndu(&mut h, unit)?);
            }
            NNR_LPS | NNR_QNT | NNR_AGG => return refuse("nnc_not_read", format!("unit type {ty}")),
            _ => return refuse("nnc_malformed", format!("unit type {ty}, which the syntax does not define")),
        }
        first = false;
        at += size;
    }
    if first {
        return refuse("nnc_truncated", "an empty bitstream");
    }
    let set = set.sorted();
    for w in set.tensors.windows(2) {
        if w[0].name == w[1].name {
            return refuse("nnc_malformed", format!("{:?} carried twice", w[0].name));
        }
    }
    Ok(set)
}

fn ndu(h: &mut Bits<'_>, unit: &[u8]) -> Result<Tensor, Refusal> {
    let pt = h.u(5)? as u32;
    let multiple = h.u(1)?;
    let fmt_present = h.u(1)?;
    let inband = h.u(1)?;
    if !matches!(pt, NNR_PT_INT | NNR_PT_RAW_FLOAT) {
        let what = match pt {
            NNR_PT_FLOAT => "NNR_PT_FLOAT (quantised parameters)".to_string(),
            NNR_PT_BLOCK => "NNR_PT_BLOCK".to_string(),
            p => format!("payload type {p}"),
        };
        return refuse("nnc_not_read", what);
    }
    if multiple != 0 {
        return refuse("nnc_not_read", "several topology elements in one data unit");
    }
    let name = h.st()?;
    if name.is_empty() {
        return refuse("nnc_malformed", "a data unit with an empty topology element name");
    }
    if pt == NNR_PT_INT && h.u(1)? != 0 {
        return refuse("nnc_not_read", format!("{name:?}: dependent quantisation (dq_flag 1)"));
    }
    if fmt_present == 1 {
        let f = h.u(7)?;
        let want = if pt == NNR_PT_INT { 0 } else { 1 };
        if f != want {
            return refuse("nnc_not_read", format!("{name:?}: decompressed data format {f}"));
        }
    }
    if inband == 0 {
        return refuse("nnc_not_read", format!("{name:?}: input parameters out of band"));
    }
    let dims_flag = h.u(1)?;
    let unary_flag = h.u(1)?;
    let cpt = h.u(4)?;
    if cpt != 0 {
        return refuse("nnc_not_read", format!("{name:?}: compressed parameter types {cpt}"));
    }
    if dims_flag == 0 {
        return refuse("nnc_not_read", format!("{name:?}: tensor dimensions out of band"));
    }
    let nd = h.ue(1)?;
    if nd == 0 || nd > 16 {
        return refuse("nnc_malformed", format!("{name:?}: {nd} dimensions"));
    }
    let mut shape = Vec::with_capacity(nd as usize);
    for _ in 0..nd {
        shape.push(h.ue(7)?);
    }
    let count = shape.iter().try_fold(1u64, |a, &d| a.checked_mul(d)).filter(|&n| n > 0);
    let Some(count) = count else { return refuse("nnc_malformed", format!("{name:?}: a dimension of 0, or a count that overflows")) };
    if count > MAX_ELEMENTS {
        return refuse("nnc_not_read", format!("{name:?}: {count} elements, past this reader's 2^28"));
    }
    let unary = if unary_flag == 1 { Some(h.u(8)? as u32) } else { None };
    if nd > 1 {
        let scan = h.u(4)?;
        if scan != 0 {
            return refuse("nnc_not_read", format!("{name:?}: scan_order {scan} (block scan with entry points)"));
        }
    }
    h.align()?;
    let start = h.bit / 8;
    let rest = &unit[start..];
    if pt == NNR_PT_RAW_FLOAT {
        if rest.len() as u64 != count * 4 {
            return refuse("nnc_malformed", format!("{name:?}: {} bytes of raw floats for {count} elements", rest.len()));
        }
        return Ok(Tensor { name, dtype: DType::F32, shape, data: rest.to_vec() });
    }
    let Some(u) = unary else { return refuse("nnc_not_read", format!("{name:?}: the unary length out of band")) };
    let (vals, used) = cabac::decode_ints(rest, count, u + 1)?;
    if used != rest.len() {
        return refuse("nnc_malformed", format!("{name:?}: the entropy-coded payload takes {used} of the unit's {} bytes", rest.len()));
    }
    let mut data = Vec::with_capacity(vals.len() * 4);
    for v in vals {
        data.extend_from_slice(&(v as i32).to_le_bytes());
    }
    Ok(Tensor { name, dtype: DType::I32, shape, data })
}

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

    fn sample() -> TensorSet {
        let mut s = TensorSet::default();
        s.metadata.insert("wai.profile".into(), "wai-int/1".into());
        s.tensors.push(Tensor::from_ints("w", DType::I32, vec![2, 3], &[1, -2, 300, -40000, 0, 7]));
        s.tensors.push(Tensor { name: "f".into(), dtype: DType::F32, shape: vec![2], data: [1.5f32.to_le_bytes(), (-0.25f32).to_le_bytes()].concat() });
        s
    }

    #[test]
    fn writes_and_reads_back() {
        let b = write(&sample()).unwrap();
        // NNR_STR first: a 15-bit size of 4, type 0, independently decodable.
        assert_eq!(&b[..4], &[0x00, 0x04, 0x02, 0x00]);
        let back = read(&b).unwrap();
        assert_eq!(back, sample().sorted());
        assert_eq!(write(&back).unwrap(), b);
    }

    #[test]
    fn narrower_integers_read_back_as_int32() {
        let mut s = TensorSet::default();
        s.tensors.push(Tensor::from_ints("x", DType::I8, vec![3], &[-128, 0, 127]));
        let back = read(&write(&s).unwrap()).unwrap();
        assert_eq!(back.tensors[0].dtype, DType::I32);
        assert_eq!(back.tensors[0].ints().unwrap(), vec![-128, 0, 127]);
    }

    #[test]
    fn refuses_outside_its_subset() {
        let b = write(&sample()).unwrap();
        let code = |x: &[u8]| read(x).err().map(|r| r.code);
        assert_eq!(code(&b[..b.len() - 1]), Some("nnc_truncated"));
        assert_eq!(code(&b[4..]), Some("nnc_malformed")); // starts at NNR_MPS
        // Turn the first data unit's payload type into NNR_PT_FLOAT.
        let mut f = b.clone();
        let mut at = 0;
        loop {
            let size = u16::from_be_bytes([f[at], f[at + 1]]) as usize;
            if f[at + 2] >> 2 == NNR_NDU as u8 {
                f[at + 3] = (f[at + 3] & 0x07) | ((NNR_PT_FLOAT as u8) << 3);
                break;
            }
            at += size;
        }
        assert_eq!(code(&f), Some("nnc_not_read"));
    }
}
