//! Tensor files as the wire forms of a prior (`extensions/prior-carriage.md` §9 to §13):
//! `wai.prior.safetensors`, `wai.prior.gguf` and the ISO/IEC 15938-17 (NNC) bitstream of
//! `wai.prior.nnc`, each with a strict reader, a canonical writer, and the integer-quantised
//! profile `wai-int/1` through which a prior carried in any of them loads, bit for bit, into
//! the native artifact an integer kernel reads.
//!
//! - [`safetensors`], [`gguf`] and [`nnc`] read a file into a [`TensorSet`] and write one
//!   canonically. A reader refuses with one of the codes of §13 ([`CODES`]).
//! - [`binding`] holds the registered bindings (§10.2): `wsr1/1` (the `WSR1` prior of
//!   exact-upscale §3.1) and `whm1/1` (the `WHM1` model of integer-payloads §9). A binding
//!   maps a native artifact to named integer tensors and back, exactly.
//! - [`load`] runs §10.3: the file in its form, canonical, in the profile, its tensors
//!   exactly the binding's, every value in range, and returns the native bytes;
//!   [`encode`] runs the other way. The digest rule (§11) is what makes the round trip one
//!   to one: a canonical file is a function of its tensors and its two metadata values.
//!
//! The corpus is `live-conformance/prior-forms/` ([`corpus`]); its second implementation is
//! `live-conformance/prior_forms_verify.py`.

use std::collections::BTreeMap;

pub mod binding;
pub mod cabac;
pub mod corpus;
pub mod gguf;
pub mod nnc;
pub mod safetensors;

/// The capability string of the safetensors form.
pub use crate::codecs::CAP_PRIOR_SAFETENSORS;
/// The capability string of the GGUF form.
pub use crate::codecs::CAP_PRIOR_GGUF;
/// The capability string of the NNC form.
pub use crate::codecs::CAP_PRIOR_NNC;

/// The integer-quantised profile's identifier (§10).
pub const PROFILE: &str = "wai-int/1";
/// The metadata key that names the profile (§10.1).
pub const KEY_PROFILE: &str = "wai.profile";
/// The metadata key that names the binding (§10.1).
pub const KEY_BINDING: &str = "wai.binding";

/// Every refusal code of §13, in the order of its table.
pub const CODES: &[&str] = &[
    "header_truncated",
    "header_too_large",
    "header_malformed",
    "bad_magic",
    "version_unsupported",
    "tensor_malformed",
    "dtype_unknown",
    "offset_mismatch",
    "header_not_canonical",
    "nnc_truncated",
    "nnc_malformed",
    "nnc_not_read",
    "profile_absent",
    "binding_unknown",
    "dtype_not_integer",
    "binding_mismatch",
    "value_out_of_range",
    "native_malformed",
];

/// Why a file was refused: a code of §13 and the particulars.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Refusal {
    pub code: &'static str,
    pub why: String,
}

impl Refusal {
    pub fn new(code: &'static str, why: impl Into<String>) -> Self {
        debug_assert!(CODES.contains(&code), "{code} is not a §13 code");
        Refusal { code, why: why.into() }
    }
}

impl std::fmt::Display for Refusal {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}: {}", self.code, self.why)
    }
}

impl std::error::Error for Refusal {}

pub(crate) fn refuse<T>(code: &'static str, why: impl Into<String>) -> Result<T, Refusal> {
    Err(Refusal::new(code, why))
}

/// An element type a reader knows (§9.3): the safetensors names, and the GGUF types they
/// correspond to, plus GGUF's block type `Q8_0`.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub enum DType {
    Bool,
    U8,
    I8,
    I16,
    U16,
    F16,
    BF16,
    I32,
    U32,
    F32,
    F64,
    I64,
    U64,
    F8E5M2,
    F8E4M3,
    /// GGUF only: blocks of 32 elements, a binary16 scale and 32 signed bytes.
    Q8_0,
}

impl DType {
    /// The safetensors name, where the format has one.
    pub fn safetensors_name(self) -> Option<&'static str> {
        Some(match self {
            DType::Bool => "BOOL",
            DType::U8 => "U8",
            DType::I8 => "I8",
            DType::I16 => "I16",
            DType::U16 => "U16",
            DType::F16 => "F16",
            DType::BF16 => "BF16",
            DType::I32 => "I32",
            DType::U32 => "U32",
            DType::F32 => "F32",
            DType::F64 => "F64",
            DType::I64 => "I64",
            DType::U64 => "U64",
            DType::F8E5M2 => "F8_E5M2",
            DType::F8E4M3 => "F8_E4M3",
            DType::Q8_0 => return None,
        })
    }
    pub fn from_safetensors_name(s: &str) -> Option<DType> {
        [
            DType::Bool, DType::U8, DType::I8, DType::I16, DType::U16, DType::F16, DType::BF16,
            DType::I32, DType::U32, DType::F32, DType::F64, DType::I64, DType::U64, DType::F8E5M2,
            DType::F8E4M3,
        ]
        .into_iter()
        .find(|d| d.safetensors_name() == Some(s))
    }
    /// The GGUF type number, where the format has one.
    pub fn gguf_type(self) -> Option<u32> {
        Some(match self {
            DType::F32 => 0,
            DType::F16 => 1,
            DType::Q8_0 => 8,
            DType::I8 => 24,
            DType::I16 => 25,
            DType::I32 => 26,
            DType::I64 => 27,
            DType::F64 => 28,
            DType::BF16 => 30,
            _ => return None,
        })
    }
    pub fn from_gguf_type(t: u32) -> Option<DType> {
        [DType::F32, DType::F16, DType::Q8_0, DType::I8, DType::I16, DType::I32, DType::I64, DType::F64, DType::BF16]
            .into_iter()
            .find(|d| d.gguf_type() == Some(t))
    }
    /// The bytes `n` elements take, or `None` when `n` is not a whole number of blocks or
    /// the count overflows.
    pub fn bytes_for(self, n: u64) -> Option<u64> {
        match self {
            DType::Bool | DType::U8 | DType::I8 | DType::F8E5M2 | DType::F8E4M3 => Some(n),
            DType::I16 | DType::U16 | DType::F16 | DType::BF16 => n.checked_mul(2),
            DType::I32 | DType::U32 | DType::F32 => n.checked_mul(4),
            DType::I64 | DType::U64 | DType::F64 => n.checked_mul(8),
            DType::Q8_0 => n.is_multiple_of(32).then(|| (n / 32).checked_mul(34)).flatten(),
        }
    }
    /// The signed integer types of the profile (§10.1): `I8`, `I16`, `I32`, `I64`.
    pub fn is_profile_integer(self) -> bool {
        matches!(self, DType::I8 | DType::I16 | DType::I32 | DType::I64)
    }
    /// The width in bytes of a profile integer type.
    pub fn int_width(self) -> Option<usize> {
        match self {
            DType::I8 => Some(1),
            DType::I16 => Some(2),
            DType::I32 => Some(4),
            DType::I64 => Some(8),
            _ => None,
        }
    }
    /// The range of a profile integer type.
    pub fn int_range(self) -> Option<(i64, i64)> {
        self.int_width().map(|w| if w == 8 { (i64::MIN, i64::MAX) } else { (-(1i64 << (8 * w - 1)), (1i64 << (8 * w - 1)) - 1) })
    }
}

/// One tensor: its name, element type, shape (outermost dimension first, the row-major
/// order of §9.1) and its bytes, little-endian.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Tensor {
    pub name: String,
    pub dtype: DType,
    pub shape: Vec<u64>,
    pub data: Vec<u8>,
}

impl Tensor {
    /// The number of elements.
    pub fn count(&self) -> Option<u64> {
        self.shape.iter().try_fold(1u64, |a, &d| a.checked_mul(d))
    }

    /// A profile integer tensor from `values`, each written at `dtype`'s width. Panics if a
    /// value does not fit: callers build tensors from values they range-checked.
    pub fn from_ints(name: &str, dtype: DType, shape: Vec<u64>, values: &[i64]) -> Tensor {
        let w = dtype.int_width().expect("a profile integer type");
        let (lo, hi) = dtype.int_range().expect("a profile integer type");
        let mut data = Vec::with_capacity(values.len() * w);
        for &v in values {
            assert!(v >= lo && v <= hi, "{name}: {v} does not fit {dtype:?}");
            data.extend_from_slice(&v.to_le_bytes()[..w]);
        }
        Tensor { name: name.to_string(), dtype, shape, data }
    }

    /// The values of a profile integer tensor, sign-extended.
    pub fn ints(&self) -> Option<Vec<i64>> {
        let w = self.dtype.int_width()?;
        if !self.data.len().is_multiple_of(w) {
            return None;
        }
        Some(
            self.data
                .chunks_exact(w)
                .map(|c| {
                    let mut b = [if c[w - 1] & 0x80 != 0 { 0xff } else { 0 }; 8];
                    b[..w].copy_from_slice(c);
                    i64::from_le_bytes(b)
                })
                .collect(),
        )
    }
}

/// A file's tensors in ascending byte order of their names, and its metadata (string
/// values only, §9.3).
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub struct TensorSet {
    pub metadata: BTreeMap<String, String>,
    pub tensors: Vec<Tensor>,
}

impl TensorSet {
    pub fn get(&self, name: &str) -> Option<&Tensor> {
        self.tensors.iter().find(|t| t.name == name)
    }
    /// Sort the tensors by name (the canonical order of §11).
    pub fn sorted(mut self) -> Self {
        self.tensors.sort_by(|a, b| a.name.as_bytes().cmp(b.name.as_bytes()));
        self
    }
}

/// A registered tensor-file form.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
pub enum Form {
    Safetensors,
    Gguf,
    Nnc,
}

impl Form {
    pub const ALL: [Form; 3] = [Form::Safetensors, Form::Gguf, Form::Nnc];
    /// The form's capability string.
    pub fn capability(self) -> &'static str {
        match self {
            Form::Safetensors => CAP_PRIOR_SAFETENSORS,
            Form::Gguf => CAP_PRIOR_GGUF,
            Form::Nnc => CAP_PRIOR_NNC,
        }
    }
    pub fn from_capability(c: &str) -> Option<Form> {
        Form::ALL.into_iter().find(|f| f.capability() == c)
    }
    /// The short name the corpus uses.
    pub fn label(self) -> &'static str {
        match self {
            Form::Safetensors => "safetensors",
            Form::Gguf => "gguf",
            Form::Nnc => "nnc",
        }
    }
}

/// Read `bytes` in `form` (the form's own rules, §9 and §12), without the canonical or
/// profile checks.
pub fn read(form: Form, bytes: &[u8]) -> Result<TensorSet, Refusal> {
    match form {
        Form::Safetensors => safetensors::read(bytes),
        Form::Gguf => gguf::read(bytes),
        Form::Nnc => nnc::read(bytes),
    }
}

/// The canonical bytes of `set` in `form` (§11).
pub fn write(form: Form, set: &TensorSet) -> Result<Vec<u8>, Refusal> {
    match form {
        Form::Safetensors => safetensors::write(set),
        Form::Gguf => gguf::write(set),
        Form::Nnc => nnc::write(set),
    }
}

/// §10.3 over a file: read it in `form`, require it canonical (§11), in the profile with a
/// registered binding (§10.1), its tensors exactly the binding's with every value in range
/// (§10.2), and return the binding's name and the native artifact's bytes.
pub fn load(form: Form, bytes: &[u8]) -> Result<(&'static str, Vec<u8>), Refusal> {
    let set = read(form, bytes)?;
    if write(form, &set)? != bytes {
        return refuse("header_not_canonical", format!("the file is not the canonical {} serialisation of its tensors (§11)", form.label()));
    }
    let b = binding::of(&set, form)?;
    let native = (b.to_native)(&set, form)?;
    Ok((b.name, native))
}

/// The other way: the native artifact `native` under `binding`, as a canonical file in
/// `form` (§10.3).
pub fn encode(binding_name: &str, native: &[u8], form: Form) -> Result<Vec<u8>, Refusal> {
    let b = binding::by_name(binding_name).ok_or_else(|| Refusal::new("binding_unknown", format!("no binding {binding_name:?}")))?;
    let mut set = (b.from_native)(native)?;
    set.metadata = binding::profile_metadata(b.name, form);
    write(form, &set.sorted())
}

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

    #[test]
    fn integer_tensors_round_trip_their_values() {
        for (d, v) in [(DType::I8, -128), (DType::I16, 32767), (DType::I32, -5), (DType::I64, i64::MIN)] {
            let t = Tensor::from_ints("x", d, vec![2], &[v, 1]);
            assert_eq!(t.ints().unwrap(), vec![v, 1]);
        }
        assert_eq!(DType::Q8_0.bytes_for(64), Some(68));
        assert_eq!(DType::Q8_0.bytes_for(33), None);
        for d in [DType::I8, DType::F8E4M3, DType::BF16, DType::U64] {
            assert_eq!(DType::from_safetensors_name(d.safetensors_name().unwrap()), Some(d));
        }
        for t in [0, 1, 8, 24, 25, 26, 27, 28, 30] {
            assert_eq!(DType::from_gguf_type(t).and_then(DType::gguf_type), Some(t));
        }
    }
}
