//! `consumption` (jwp-receipts §2.4.13; `extensions/ai-usage.md` §6): an agent's statement of
//! the objects it consumed, the use it made of them in each AI usage vocabulary it answers,
//! the usage statements it saw and the room grant it held.
//!
//! This file holds steps 2 and 4 of jwp-receipts §3.1: the body's members and their types,
//! and the class rules of ai-usage §6.2. Step 6, the check against the objects, the
//! statements and the room, is `crate::ai_usage`'s, which holds them.

use super::{get_opt_hash, hex, members, unhex, Reason};
use serde_json::{Map, Value};

/// The body's members, in the order step 2 reads them (ai-usage §6.1).
const CONSUMPTION_MEMBERS: &[&str] = &["objects", "usage", "seen", "grant"];
/// A `usage` entry's members.
const USAGE_MEMBERS: &[&str] = &["vocabulary", "term"];
/// The most objects one claim names.
///
/// cbindgen:ignore
pub const MAX_OBJECTS: usize = 1024;
/// The most statements one claim says it saw.
///
/// cbindgen:ignore
pub const MAX_SEEN: usize = 64;
/// The longest term, in bytes.
///
/// cbindgen:ignore
pub const MAX_TERM: usize = 255;

/// The three vocabularies a usage statement carries (ai-usage §3.1), in the order a claim
/// lists them: ascending by name.
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub enum Vocabulary {
    /// `draft-ietf-aipref-vocab-08`.
    Aipref,
    /// The Training and Data Mining Assertion, version 1.1.
    Cawg,
    /// The TDM Reservation Protocol.
    Tdmrep,
}

impl Vocabulary {
    /// cbindgen:ignore
    pub const ALL: [Vocabulary; 3] = [Vocabulary::Aipref, Vocabulary::Cawg, Vocabulary::Tdmrep];

    pub fn label(self) -> &'static str {
        match self {
            Vocabulary::Aipref => "aipref",
            Vocabulary::Cawg => "cawg",
            Vocabulary::Tdmrep => "tdmrep",
        }
    }

    pub fn from_label(s: &str) -> Option<Vocabulary> {
        Vocabulary::ALL.into_iter().find(|v| v.label() == s)
    }

    /// Whether `term` is a term of this vocabulary (ai-usage §6.2 rule 4).
    pub fn has_term(self, term: &str) -> bool {
        match self {
            Vocabulary::Aipref => AIPREF_LABELS.contains(&term),
            Vocabulary::Tdmrep => term == TDM_TERM,
            Vocabulary::Cawg => CAWG_KEYS.contains(&term) || (cawg_key_form(term) && !term.starts_with("cawg.")),
        }
    }
}

/// The usage category labels of `draft-ietf-aipref-vocab-08` §6.1, Table 1.
pub const AIPREF_LABELS: [&str; 3] = ["train-ai", "ai-use", "search"];
/// The one term of the TDM Reservation Protocol: the property that governs mining.
pub const TDM_TERM: &str = "tdm-reservation";
/// The four entries the Training and Data Mining Assertion 1.1 defines (§3.1 there).
pub const CAWG_KEYS: [&str; 4] = ["cawg.data_mining", "cawg.ai_inference", "cawg.ai_generative_training", "cawg.ai_training"];

/// ai-usage §3.4: 1 to 255 bytes of components separated by single dots, each one or more
/// of `a`–`z`, `A`–`Z`, `0`–`9`, `_` and `-`.
pub fn cawg_key_form(k: &str) -> bool {
    !k.is_empty()
        && k.len() <= 255
        && k.split('.').all(|c| !c.is_empty() && c.bytes().all(|b| b.is_ascii_alphanumeric() || b == b'_' || b == b'-'))
}

/// One `usage` entry: the use, as a vocabulary names it, or `None` for a use outside it.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct UsageEntry {
    pub vocabulary: Vocabulary,
    pub term: Option<String>,
}

/// A `consumption` body.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct ConsumptionClaim {
    /// The content hashes of the objects consumed.
    pub objects: Vec<[u8; 32]>,
    /// The use, in each vocabulary the agent answers.
    pub usage: Vec<UsageEntry>,
    /// The digests of the statements that governed the objects.
    pub seen: Vec<[u8; 32]>,
    /// The digest of the room grant the agent consumed under.
    pub grant: Option<[u8; 32]>,
}

fn strictly_ascending<T: Ord>(v: &[T]) -> bool {
    v.windows(2).all(|w| w[0] < w[1])
}

impl ConsumptionClaim {
    pub(super) fn to_value(&self) -> Value {
        let mut m = Map::new();
        m.insert("objects".into(), Value::Array(self.objects.iter().map(|h| hex(h).into()).collect()));
        m.insert(
            "usage".into(),
            Value::Array(
                self.usage
                    .iter()
                    .map(|u| {
                        let mut e = Map::new();
                        e.insert("vocabulary".into(), u.vocabulary.label().into());
                        e.insert("term".into(), u.term.clone().map_or(Value::Null, Value::String));
                        Value::Object(e)
                    })
                    .collect(),
            ),
        );
        m.insert("seen".into(), Value::Array(self.seen.iter().map(|h| hex(h).into()).collect()));
        m.insert("grant".into(), self.grant.map_or(Value::Null, |g| hex(&g).into()));
        Value::Object(m)
    }

    /// Step 2: the members — missing, then unknown — then each member's type in the table's
    /// order, an entry's members and types where the table reaches it.
    pub(super) fn read(m: &Map<String, Value>) -> Result<ConsumptionClaim, Reason> {
        members(m, CONSUMPTION_MEMBERS)?;
        let ft = Reason::FieldType;
        let hashes = |k: &str, least: usize, most: usize| -> Result<Vec<[u8; 32]>, Reason> {
            let a = m[k].as_array().ok_or(ft)?;
            if a.len() < least || a.len() > most {
                return Err(ft);
            }
            a.iter().map(|v| v.as_str().and_then(unhex).and_then(|b| b.try_into().ok()).ok_or(ft)).collect()
        };
        let objects = hashes("objects", 1, MAX_OBJECTS)?;
        let entries = m["usage"].as_array().ok_or(ft)?;
        if entries.is_empty() || entries.len() > Vocabulary::ALL.len() {
            return Err(ft);
        }
        let usage = entries
            .iter()
            .map(|v| {
                let e = v.as_object().ok_or(ft)?;
                members(e, USAGE_MEMBERS)?;
                let vocabulary = Vocabulary::from_label(e["vocabulary"].as_str().ok_or(ft)?).ok_or(ft)?;
                let term = match &e["term"] {
                    Value::Null => None,
                    Value::String(s) if !s.is_empty() && s.len() <= MAX_TERM => Some(s.clone()),
                    _ => return Err(ft),
                };
                Ok(UsageEntry { vocabulary, term })
            })
            .collect::<Result<Vec<_>, Reason>>()?;
        let seen = hashes("seen", 0, MAX_SEEN)?;
        let grant = get_opt_hash(m, "grant")?;
        Ok(ConsumptionClaim { objects, usage, seen, grant })
    }

    /// The counts and term lengths step 2 checks, for a body built field by field.
    pub(super) fn check_values(&self) -> Result<(), Reason> {
        let ft = Reason::FieldType;
        if self.objects.is_empty() || self.objects.len() > MAX_OBJECTS || self.seen.len() > MAX_SEEN {
            return Err(ft);
        }
        if self.usage.is_empty() || self.usage.len() > Vocabulary::ALL.len() {
            return Err(ft);
        }
        if self.usage.iter().any(|u| u.term.as_ref().is_some_and(|t| t.is_empty() || t.len() > MAX_TERM)) {
            return Err(ft);
        }
        Ok(())
    }

    /// Step 4: ai-usage §6.2's rules, in its order.
    pub(super) fn rules(&self, subject: &[u8; 32], signer_pubkey: &[u8; 32], signer_id: &str) -> Result<(), Reason> {
        if !strictly_ascending(&self.objects) {
            return Err(Reason::ConsumptionObjectsUnordered);
        }
        if self.objects.first() != Some(subject) {
            return Err(Reason::ConsumptionSubjectMismatch);
        }
        let vocabularies: Vec<Vocabulary> = self.usage.iter().map(|u| u.vocabulary).collect();
        if !strictly_ascending(&vocabularies) {
            return Err(Reason::ConsumptionUsageUnordered);
        }
        if self.usage.iter().any(|u| u.term.as_ref().is_some_and(|t| !u.vocabulary.has_term(t))) {
            return Err(Reason::ConsumptionTermUnknown);
        }
        if !strictly_ascending(&self.seen) {
            return Err(Reason::ConsumptionSeenUnordered);
        }
        if crate::signer_id::did_key(signer_pubkey) != signer_id {
            return Err(Reason::ConsumptionSignerNotPrincipal);
        }
        Ok(())
    }
}

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

    #[test]
    fn terms_are_the_vocabularies_own() {
        for t in AIPREF_LABELS {
            assert!(Vocabulary::Aipref.has_term(t));
        }
        assert!(!Vocabulary::Aipref.has_term("inference"));
        assert!(!Vocabulary::Aipref.has_term("train-genai"));
        assert!(Vocabulary::Tdmrep.has_term("tdm-reservation"));
        assert!(!Vocabulary::Tdmrep.has_term("tdm-policy"));
        for k in CAWG_KEYS {
            assert!(Vocabulary::Cawg.has_term(k));
        }
        assert!(Vocabulary::Cawg.has_term("org.example.summaries"));
        assert!(!Vocabulary::Cawg.has_term("cawg.ai_summaries"), "the prefix is reserved");
        assert!(!Vocabulary::Cawg.has_term("org..example"));
        assert!(!Vocabulary::Cawg.has_term("org example"));
        assert!(!Vocabulary::Cawg.has_term(""));
        assert!(!Vocabulary::Cawg.has_term(&"a".repeat(256)));
        assert!(Vocabulary::Cawg.has_term(&"a".repeat(255)));
    }

    #[test]
    fn vocabularies_order_by_name() {
        let mut v = Vocabulary::ALL.map(|x| x.label());
        v.sort();
        assert_eq!(v, Vocabulary::ALL.map(|x| x.label()));
    }
}
