//! Counted bytes of an object (`extensions/staged-measurement.md` §3).
//!
//! The counted bytes of a `WAI1` envelope are `12 + man_len + P`, and of a `WAI2`
//! envelope `8 + man_len + 2 + 8·n + Σ Pᵢ` (SPEC §2, §7.1): the envelope exactly as
//! written, which a well-formed envelope is to the byte. An object is countable only
//! where WAI checks the format of every payload it carries (§3.2): a capability SPEC §5
//! pins to a binary format (`WIV1`, `WIH1`, `WIA1`, `WIS1`, and `WIR1` for a
//! refinement-only capability), whose payload parses as that format by the parse a sink
//! runs (`container::check_int_payload_framing`), and, where the manifest names a
//! fallback with a pinned format, parses as the fallback's too (as
//! `Wai::check_int_payload_framing` checks a payload a fallback decodes). Every other
//! capability, a classical codestream or a string WAI does not register among them, is
//! not countable (`no-format-check`): nothing in WAI tells its bytes from a text
//! encoding of them. A `WAI2` component whose `sha256` does not match its payload is not
//! countable (`sha256-mismatch`), and neither is an encrypted payload (`encrypted`),
//! whose plaintext's format cannot be checked.
//!
//! A countable object's figure carries the framing label `registered` (§3.3). Nothing
//! here estimates a figure for an object that is not countable.

use crate::container::{check_int_payload_framing, manifest_v2_security, Wai, WaiMulti, MAGIC, MAGIC_V2};

/// The framing label of a figure whose every payload is in its capability's registered
/// format (§3.3).
pub const REGISTERED: &str = "registered";

/// An object's counted bytes, split as §3.4 states them.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Counted {
    /// The envelope's own structure: magic and length prefixes, and for `WAI2` the
    /// rendition count and table.
    pub header: u64,
    /// The manifest's bytes as written.
    pub manifest: u64,
    /// The payloads' bytes.
    pub payload: u64,
    /// `header + manifest + payload`: the envelope's length.
    pub total: u64,
    /// The framing label (§3.3).
    pub framing: &'static str,
}

/// Why an object has no counted bytes.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum FramingError {
    /// The bytes are not a `WAI1` or `WAI2` envelope.
    Envelope(String),
    /// Bytes follow the envelope's last payload: the object is not the envelope its
    /// lengths describe.
    Trailing { envelope: u64, object: u64 },
    /// A payload is not countable (§3.2): `component` is the `WAI2` entry's index,
    /// `None` for a `WAI1`; `capability` the capability whose check refused it; `code`
    /// `no-format-check`, `sha256-mismatch`, `encrypted`, or the format parse's code
    /// (`bad_magic`, `format_capability_mismatch`, … integer-payloads §8, staged-delivery
    /// A.6).
    NotCountable { component: Option<usize>, capability: String, code: &'static str, detail: String },
}

impl FramingError {
    /// A short token for the error: `not-an-envelope`, `trailing`, or the code of a
    /// payload that is not countable.
    pub fn code(&self) -> &'static str {
        match self {
            Self::Envelope(_) => "not-an-envelope",
            Self::Trailing { .. } => "trailing",
            Self::NotCountable { code, .. } => code,
        }
    }
}

impl std::fmt::Display for FramingError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Envelope(e) => write!(f, "not an envelope: {e}"),
            Self::Trailing { envelope, object } => write!(f, "{} bytes follow the envelope's {envelope}", object - envelope),
            Self::NotCountable { component: Some(i), capability, code, detail } => write!(f, "component {i} ({capability}) is not countable: {code}: {detail}"),
            Self::NotCountable { component: None, capability, code, detail } => write!(f, "the payload ({capability}) is not countable: {code}: {detail}"),
        }
    }
}

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

/// Whether WAI checks the format of `capability`'s payload: a pinned integer format, or
/// `WIR1` for a refinement-only capability.
pub fn has_format_check(capability: &str) -> bool {
    crate::codecs::pinned_int_payload(capability).is_some() || crate::codecs::is_refinement_only(capability)
}

/// Why `payload` is not countable under `capability` (its code and what failed), or
/// `None` when it is (§3.2).
pub fn not_countable(capability: &str, payload: &[u8]) -> Option<(&'static str, String)> {
    if !has_format_check(capability) {
        return Some(("no-format-check", "WAI checks no format for this capability's payload".into()));
    }
    check_int_payload_framing(capability, payload).err().map(|e| (e.code, e.detail.to_string()))
}

fn refuse(component: Option<usize>, capability: &str, (code, detail): (&'static str, String)) -> FramingError {
    FramingError::NotCountable { component, capability: capability.to_string(), code, detail }
}

/// The counted bytes of one object (§3.1).
pub fn count_object(envelope: &[u8]) -> Result<Counted, FramingError> {
    let len = envelope.len() as u64;
    let magic = envelope.get(..4).ok_or_else(|| FramingError::Envelope("shorter than a magic".into()))?;
    let encrypted = || ("encrypted", "an encrypted payload's format cannot be checked".to_string());
    let (header, manifest, payload) = if magic == MAGIC {
        let w = Wai::from_bytes(envelope).map_err(|e| FramingError::Envelope(e.to_string()))?;
        let man = u32::from_le_bytes(envelope[4..8].try_into().expect("checked by the parse")) as u64;
        let cap = w.capability();
        if w.manifest.security.as_ref().is_some_and(|s| !s.is_plaintext()) {
            return Err(refuse(None, cap, encrypted()));
        }
        if let Some(why) = not_countable(cap, &w.payload) {
            return Err(refuse(None, cap, why));
        }
        // The fallback decodes the same payload (SPEC §4 step 5), so a fallback with a
        // pinned format must find it in that format too.
        if let Some(fb) = w.fallback().filter(|fb| has_format_check(fb)) {
            if let Some(why) = not_countable(fb, &w.payload) {
                return Err(refuse(None, fb, why));
            }
        }
        (12, man, w.payload.len() as u64)
    } else if magic == MAGIC_V2 {
        let w = WaiMulti::from_bytes(envelope).map_err(|e| FramingError::Envelope(e.to_string()))?;
        let man = u32::from_le_bytes(envelope[4..8].try_into().expect("checked by the parse")) as u64;
        if manifest_v2_security(&w.manifest).map_err(FramingError::Envelope)?.is_some_and(|s| !s.is_plaintext()) {
            return Err(refuse(None, "", encrypted()));
        }
        let mut payload = 0u64;
        for (i, (meta, p)) in w.manifest.renditions.iter().zip(&w.renditions).enumerate() {
            if meta.sha256.is_some() && !meta.payload_digest_matches(p) {
                return Err(refuse(Some(i), &meta.capability, ("sha256-mismatch", "the entry's sha256 is not its payload's".into())));
            }
            if let Some(why) = not_countable(&meta.capability, p) {
                return Err(refuse(Some(i), &meta.capability, why));
            }
            payload += p.len() as u64;
        }
        (8 + 2 + 8 * w.renditions.len() as u64, man, payload)
    } else {
        return Err(FramingError::Envelope("neither WAI1 nor WAI2".into()));
    };
    let total = header + manifest + payload;
    if total != len {
        return Err(FramingError::Trailing { envelope: total, object: len });
    }
    Ok(Counted { header, manifest, payload, total, framing: REGISTERED })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::container::{Conditioning, Manifest, ModelRequirement, Rendition, Security};

    const WIH1: &[u8] = include_bytes!("../tests/vectors/int_hyper/tiny_bitstream.wih1");
    const HYPER: &str = "wai.neural.int_hyper";

    fn wai1(cap: &str, fallback: Option<&str>, payload: Vec<u8>) -> Vec<u8> {
        let mut mr = ModelRequirement::new(cap);
        if let Some(f) = fallback {
            mr = mr.with_fallback(f);
        }
        Wai::new(Manifest::new("image", "replicate", mr, Conditioning::new("k")), payload).to_bytes().unwrap()
    }

    fn code(r: Result<Counted, FramingError>) -> &'static str {
        r.err().map(|e| e.code()).unwrap_or("counted")
    }

    #[test]
    fn a_wai1_counts_its_bytes_split_three_ways() {
        let e = wai1(HYPER, None, WIH1.to_vec());
        let c = count_object(&e).unwrap();
        let man = u32::from_le_bytes(e[4..8].try_into().unwrap()) as u64;
        assert_eq!(c, Counted { header: 12, manifest: man, payload: WIH1.len() as u64, total: e.len() as u64, framing: REGISTERED });
    }

    #[test]
    fn a_wai2_counts_its_table_and_checks_its_digests() {
        let two = || {
            WaiMulti::new("image", "replicate", serde_json::Value::Null, vec![
                Rendition::media(HYPER, "int_hyper", WIH1.to_vec()).with_payload_digest(),
                Rendition::media(HYPER, "int_hyper", WIH1.to_vec()),
            ])
        };
        let e = two().to_bytes().unwrap();
        let c = count_object(&e).unwrap();
        assert_eq!((c.header, c.payload, c.total), (8 + 2 + 16, 2 * WIH1.len() as u64, e.len() as u64));
        // A stated digest that is not the payload's.
        let mut bad = two();
        bad.manifest.renditions[0].sha256 = Some("0".repeat(64));
        assert!(matches!(count_object(&bad.to_bytes().unwrap()), Err(FramingError::NotCountable { component: Some(0), code: "sha256-mismatch", .. })));
    }

    #[test]
    fn only_a_checked_format_is_counted() {
        // A capability WAI checks no format for, a classical codestream or a string WAI
        // does not register, is not countable.
        assert_eq!(code(count_object(&wai1("wai.image.png", None, b"\x89PNG\r\n\x1a\n".to_vec()))), "no-format-check");
        assert_eq!(code(count_object(&wai1("wai.image.unregistered", None, vec![1, 2, 3]))), "no-format-check");
        // Text forms of a pinned format, and bytes before it, fail its parse.
        let b64: Vec<u8> = WIH1.iter().map(|b| b"ABCDEFGHIJKLMNOP"[(b & 15) as usize]).collect();
        for p in [b"{\"rans_b64\": \"AAAA\"}".to_vec(), b"\"V0lIMQE=\"".to_vec(), b64, [b"\xef\xbb\xbf".as_slice(), WIH1].concat()] {
            assert_eq!(code(count_object(&wai1(HYPER, None, p))), "bad_magic");
        }
        // The fallback decodes the same payload: a WIH1 payload under an int_motion
        // fallback is not that fallback's format.
        assert_eq!(code(count_object(&wai1(HYPER, Some("wai.video.int_motion"), WIH1.to_vec()))), "format_capability_mismatch");
        // A fallback WAI checks no format for is not checked.
        assert_eq!(code(count_object(&wai1(HYPER, Some("wai.image.png"), WIH1.to_vec()))), "counted");
    }

    #[test]
    fn trailing_truncated_and_encrypted_are_refused() {
        let mut e = wai1(HYPER, None, WIH1.to_vec());
        e.push(0);
        assert!(matches!(count_object(&e), Err(FramingError::Trailing { .. })));
        e.truncate(e.len() - 3);
        assert!(matches!(count_object(&e), Err(FramingError::Envelope(_))));
        assert!(matches!(count_object(b"WAI"), Err(FramingError::Envelope(_))));
        let mut m = Manifest::new("image", "replicate", ModelRequirement::new(HYPER), Conditioning::new("k"));
        m.security = Some(Security::new("aes-256-gcm", 1));
        let enc = Wai::new(m, WIH1.to_vec()).to_bytes().unwrap();
        assert_eq!(code(count_object(&enc)), "encrypted");
    }
}
