//! Sealed sets: `wai.prior.sealed` (extensions/prior-carriage.md §18.2).
//!
//! A one-file set encrypted under a 16-byte AES-128 key that its origin releases to the sinks it
//! entitles. The pin covers the sealed bytes, so any server may serve them and a sink checks them
//! by digest before it holds a key; opened, the plaintext is checked against the SHA-256 the
//! header states. A sink without the key holds bytes it cannot read, which SPEC §3.1 check 6
//! does not hold (`prior_not_held`), and reports why (`set_key_absent`).

use super::usage::{self, Usage};
use crate::codecs::{CAP_PRIOR_GGUF, CAP_PRIOR_SAFETENSORS};
use crate::pin::sha256;
use aes_gcm::aes::cipher::{generic_array::GenericArray, BlockEncrypt, KeyInit};
use aes_gcm::aes::Aes128;

/// The magic a sealed set begins with.
pub const MAGIC: [u8; 4] = *b"WSK1";
/// The longest plaintext: 2^40 bytes.
pub const MAX_PLAINTEXT: u64 = 1 << 40;

pub const SEALED_MALFORMED: &str = "sealed_malformed";
pub const SEALED_INNER_FORM: &str = "sealed_inner_form";
pub const SET_KEY_ABSENT: &str = "set_key_absent";
pub const SET_KEY_MISMATCH: &str = "set_key_mismatch";

/// The inner forms a sealed set admits: a capability's one file (no form), and the two
/// tensor-file forms, each established exact (§5).
pub fn inner_form_admitted(form: &str) -> bool {
    form.is_empty() || form == CAP_PRIOR_SAFETENSORS || form == CAP_PRIOR_GGUF
}

/// A sealed set's header, read without its key.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Header {
    /// The inner form; empty for a capability's one file.
    pub inner_form: String,
    pub kid: [u8; 16],
    pub iv: [u8; 16],
    /// The SHA-256 of the plaintext: the set's digest in its inner form.
    pub inner: [u8; 32],
    pub usage: Option<Usage>,
    /// Where the ciphertext begins, and its length.
    pub at: usize,
    pub len: u64,
}

/// AES-128 in counter mode (NIST SP 800-38A §6.5): the counter block `iv`, incremented as one
/// 128-bit big-endian integer per 16-byte block. Encryption and decryption are the same.
pub fn ctr(key: &[u8; 16], iv: &[u8; 16], data: &[u8]) -> Vec<u8> {
    let cipher = Aes128::new(GenericArray::from_slice(key));
    let mut counter = u128::from_be_bytes(*iv);
    let mut out = Vec::with_capacity(data.len());
    for chunk in data.chunks(16) {
        let mut block = GenericArray::clone_from_slice(&counter.to_be_bytes());
        cipher.encrypt_block(&mut block);
        out.extend(chunk.iter().zip(block.iter()).map(|(a, b)| a ^ b));
        counter = counter.wrapping_add(1);
    }
    out
}

fn le(b: &[u8]) -> u64 {
    b.iter().rev().fold(0u64, |a, x| a << 8 | *x as u64)
}

/// §18.2's header checks, in order: the layout (`sealed_malformed`), the inner form
/// (`sealed_inner_form`), the usage statement (`set_usage_malformed`).
pub fn read_header(b: &[u8]) -> Result<Header, &'static str> {
    let bad = SEALED_MALFORMED;
    if b.len() < 5 || b[..4] != MAGIC {
        return Err(bad);
    }
    let k = b[4] as usize;
    if k > 64 || b.len() < 73 + k {
        return Err(bad);
    }
    let form = &b[5..5 + k];
    if !form.iter().all(|c| (0x21..=0x7e).contains(c)) {
        return Err(bad);
    }
    let inner_form = String::from_utf8(form.to_vec()).map_err(|_| bad)?;
    let kid: [u8; 16] = b[5 + k..21 + k].try_into().map_err(|_| bad)?;
    let iv: [u8; 16] = b[21 + k..37 + k].try_into().map_err(|_| bad)?;
    let inner: [u8; 32] = b[37 + k..69 + k].try_into().map_err(|_| bad)?;
    let u = le(&b[69 + k..73 + k]) as usize;
    if u > usage::MAX_USAGE || b.len() < 81 + k + u {
        return Err(bad);
    }
    let len = le(&b[73 + k + u..81 + k + u]);
    let at = 81 + k + u;
    if len == 0 || len > MAX_PLAINTEXT || (b.len() - at) as u64 != len {
        return Err(bad);
    }
    if !inner_form_admitted(&inner_form) {
        return Err(SEALED_INNER_FORM);
    }
    let usage = if u == 0 { None } else { Some(usage::read(&b[73 + k..73 + k + u])?) };
    Ok(Header { inner_form, kid, iv, inner, usage, at, len })
}

/// Opens a sealed set with the key a sink holds for its `kid`, if any: the header and the
/// plaintext, checked against the header's SHA-256; or the reason (`set_key_absent`,
/// `set_key_mismatch`, or a header's).
pub fn open(b: &[u8], key_for: &dyn Fn(&[u8; 16]) -> Option<[u8; 16]>) -> Result<(Header, Vec<u8>), &'static str> {
    let h = read_header(b)?;
    let key = key_for(&h.kid).ok_or(SET_KEY_ABSENT)?;
    let plain = ctr(&key, &h.iv, &b[h.at..]);
    if sha256(&plain) != h.inner {
        return Err(SET_KEY_MISMATCH);
    }
    Ok((h, plain))
}

/// Seals `plaintext`, a set in `inner_form`, under `key` (named `kid`) from the counter block
/// `iv`, with the usage statement `usage` (its JSON text) or none.
pub fn seal(plaintext: &[u8], inner_form: &str, kid: &[u8; 16], key: &[u8; 16], iv: &[u8; 16], usage: Option<&[u8]>) -> Vec<u8> {
    let u = usage.unwrap_or(&[]);
    let mut out = MAGIC.to_vec();
    out.push(inner_form.len() as u8);
    out.extend_from_slice(inner_form.as_bytes());
    out.extend_from_slice(kid);
    out.extend_from_slice(iv);
    out.extend_from_slice(&sha256(plaintext));
    out.extend_from_slice(&(u.len() as u32).to_le_bytes());
    out.extend_from_slice(u);
    out.extend_from_slice(&(plaintext.len() as u64).to_le_bytes());
    out.extend_from_slice(&ctr(key, iv, plaintext));
    out
}

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

    /// NIST SP 800-38A F.5.1, CTR-AES128.Encrypt, its first two blocks.
    #[test]
    fn counter_mode_matches_sp_800_38a() {
        let h = |s: &str| (0..s.len()).step_by(2).map(|i| u8::from_str_radix(&s[i..i + 2], 16).expect("hex")).collect::<Vec<u8>>();
        let key: [u8; 16] = h("2b7e151628aed2a6abf7158809cf4f3c").try_into().expect("16");
        let iv: [u8; 16] = h("f0f1f2f3f4f5f6f7f8f9fafbfcfdfeff").try_into().expect("16");
        let plain = h("6bc1bee22e409f96e93d7e117393172aae2d8a571e03ac9c9eb76fac45af8e51");
        assert_eq!(ctr(&key, &iv, &plain), h("874d6191b620e3261bef6864990db6ce9806f66b7970fdff8617187bb9fffdff"));
    }

    #[test]
    fn sealed_and_opened() {
        let plain = b"a one-file parameter set".to_vec();
        let (kid, key, iv) = ([1u8; 16], [2u8; 16], [3u8; 16]);
        let s = seal(&plain, "", &kid, &key, &iv, None);
        let (h, got) = open(&s, &|k| (*k == kid).then_some(key)).expect("opens");
        assert_eq!((got, h.inner_form.as_str()), (plain.clone(), ""));
        assert_eq!(open(&s, &|_| None).map(|_| ()), Err(SET_KEY_ABSENT));
        assert_eq!(open(&s, &|_| Some([9u8; 16])).map(|_| ()), Err(SET_KEY_MISMATCH));
        for inner in ["wai.prior.nnc", "wai.prior.set", "wai.prior.sealed", "wai.prior.tree", "x.unregistered"] {
            assert_eq!(read_header(&seal(&plain, inner, &kid, &key, &iv, None)), Err(SEALED_INNER_FORM), "{inner}");
        }
        assert_eq!(read_header(&s[..s.len() - 1]), Err(SEALED_MALFORMED));
    }
}
