//! Signed vault manifest (docs/03 "Integrity and rollback protection").
//!
//! Body: canonical CBOR (RFC 8949 §4.2.1 deterministic encoding) of a map with text keys:
//!
//! | key            | type        | meaning                                                   |
//! |----------------|-------------|-----------------------------------------------------------|
//! | `prev`         | bstr (32)   | blake3 of the previous manifest body, zeros for version 1 |
//! | `items`        | bstr (32)   | Merkle root (below)                                       |
//! | `version`      | uint        | +1 per publish                                            |
//! | `vault_id`     | bstr (16)   |                                                           |
//! | `device_id`    | bstr (16)   | publishing device                                         |
//! | `created_at`   | uint        | unix seconds                                              |
//! | `head_revision`| uint        | server revision the Merkle root describes                 |
//!
//! Merkle tree over every item the server holds (tombstones included), sorted by `item_id` bytes:
//! * `content = blake3(cbor([envelope(wrapped_item_key), envelope(payload)]))`, or `blake3(cbor([]))` for a tombstone,
//!   where `envelope(e)` is the canonical CBOR map `{n, v, ct, alg, enc, kid}` (`n`/`enc` omitted when empty).
//! * `leaf = blake3(0x00 || item_id[16] || revision as u64 big-endian || content)`
//! * `node = blake3(0x01 || left || right)`; an odd node at the end of a level is promoted unchanged.
//! * The root of an empty vault is 32 zero bytes.
//!
//! `signature = Ed25519(account sig_sk, body)`, `manifest hash = blake3(body)`.

use kv_crypto::keys::{verify_signature, AccountKeyset};
use kv_proto::{Envelope, ItemRecord, Manifest};
use uuid::Uuid;

#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ManifestBody {
    pub vault_id: Uuid,
    pub version: u64,
    pub head_revision: u64,
    pub items: [u8; 32],
    pub prev: [u8; 32],
    pub device_id: Uuid,
    pub created_at: u64,
}

/// Why a manifest was rejected. Every variant is an integrity alarm.
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum ManifestError {
    #[error("manifest is not canonical CBOR or has the wrong shape")]
    Malformed,
    #[error("manifest signature is invalid")]
    BadSignature,
    #[error("manifest belongs to another vault")]
    WrongVault,
    #[error("manifest version {got} is older than the last seen {seen} (rollback)")]
    Rollback { seen: u64, got: u64 },
    #[error("manifest version {0} differs from the one seen before (fork)")]
    Fork(u64),
    #[error("manifest prev hash does not chain to the last seen manifest")]
    ChainBroken,
    #[error("manifest Merkle root does not match the items the server sent")]
    MerkleMismatch,
    #[error("server has no manifest but one was seen before")]
    Missing,
}

// ---------------------------------------------------------------- minimal canonical CBOR writer

mod cbor {
    pub fn head(out: &mut Vec<u8>, major: u8, n: u64) {
        let m = major << 5;
        match n {
            0..=23 => out.push(m | n as u8),
            24..=0xff => out.extend_from_slice(&[m | 24, n as u8]),
            0x100..=0xffff => {
                out.push(m | 25);
                out.extend_from_slice(&(n as u16).to_be_bytes())
            }
            0x1_0000..=0xffff_ffff => {
                out.push(m | 26);
                out.extend_from_slice(&(n as u32).to_be_bytes())
            }
            _ => {
                out.push(m | 27);
                out.extend_from_slice(&n.to_be_bytes())
            }
        }
    }
    pub fn uint(out: &mut Vec<u8>, n: u64) {
        head(out, 0, n)
    }
    pub fn bytes(out: &mut Vec<u8>, b: &[u8]) {
        head(out, 2, b.len() as u64);
        out.extend_from_slice(b);
    }
    pub fn text(out: &mut Vec<u8>, s: &str) {
        head(out, 3, s.len() as u64);
        out.extend_from_slice(s.as_bytes());
    }
    pub fn array(out: &mut Vec<u8>, len: usize) {
        head(out, 4, len as u64)
    }
    pub fn map(out: &mut Vec<u8>, len: usize) {
        head(out, 5, len as u64)
    }
}

/// Canonical CBOR of an envelope (keys in deterministic order: n, v, ct, alg, enc, kid).
pub fn envelope_cbor(e: &Envelope) -> Vec<u8> {
    let mut out = Vec::with_capacity(e.ct.len() + 96);
    let alg = serde_json::to_value(e.alg).ok().and_then(|v| v.as_str().map(str::to_string)).unwrap_or_default();
    let n = 4 + usize::from(!e.n.is_empty()) + usize::from(!e.enc.is_empty());
    cbor::map(&mut out, n);
    if !e.n.is_empty() {
        cbor::text(&mut out, "n");
        cbor::bytes(&mut out, &e.n);
    }
    cbor::text(&mut out, "v");
    cbor::uint(&mut out, e.v.into());
    cbor::text(&mut out, "ct");
    cbor::bytes(&mut out, &e.ct);
    cbor::text(&mut out, "alg");
    cbor::text(&mut out, &alg);
    if !e.enc.is_empty() {
        cbor::text(&mut out, "enc");
        cbor::bytes(&mut out, &e.enc);
    }
    cbor::text(&mut out, "kid");
    cbor::bytes(&mut out, e.kid.as_bytes());
    out
}

/// `blake3(envelope)` for one item record. A sealed tombstone hashes its envelope like a live record; only
/// a legacy tombstone without ciphertext hashes the empty array.
pub fn content_hash(record: &ItemRecord) -> [u8; 32] {
    let mut out = Vec::new();
    match (&record.wrapped_item_key, &record.payload) {
        (Some(k), Some(p)) => {
            cbor::array(&mut out, 2);
            out.extend_from_slice(&envelope_cbor(k));
            out.extend_from_slice(&envelope_cbor(p));
        }
        _ => cbor::array(&mut out, 0),
    }
    *blake3::hash(&out).as_bytes()
}

/// Merkle root over `(item_id, revision, content_hash)`; sorts its input.
pub fn merkle_root(mut leaves: Vec<(Uuid, u64, [u8; 32])>) -> [u8; 32] {
    leaves.sort_by(|a, b| a.0.as_bytes().cmp(b.0.as_bytes()));
    let mut level: Vec<[u8; 32]> = leaves
        .iter()
        .map(|(id, rev, h)| {
            let mut hs = blake3::Hasher::new();
            hs.update(&[0]);
            hs.update(id.as_bytes());
            hs.update(&rev.to_be_bytes());
            hs.update(h);
            *hs.finalize().as_bytes()
        })
        .collect();
    if level.is_empty() {
        return [0; 32];
    }
    while level.len() > 1 {
        level = level
            .chunks(2)
            .map(|pair| match pair {
                [l, r] => {
                    let mut hs = blake3::Hasher::new();
                    hs.update(&[1]);
                    hs.update(l);
                    hs.update(r);
                    *hs.finalize().as_bytes()
                }
                [single] => *single,
                _ => unreachable!(),
            })
            .collect();
    }
    level[0]
}

/// Merkle root of the server's view of a vault, from records with a server revision.
pub fn merkle_root_of<'a>(records: impl IntoIterator<Item = &'a ItemRecord>) -> [u8; 32] {
    merkle_root(records.into_iter().map(|r| (r.item_id, r.revision, content_hash(r))).collect())
}

impl ManifestBody {
    pub fn encode(&self) -> Vec<u8> {
        let mut out = Vec::with_capacity(200);
        cbor::map(&mut out, 7);
        cbor::text(&mut out, "prev");
        cbor::bytes(&mut out, &self.prev);
        cbor::text(&mut out, "items");
        cbor::bytes(&mut out, &self.items);
        cbor::text(&mut out, "version");
        cbor::uint(&mut out, self.version);
        cbor::text(&mut out, "vault_id");
        cbor::bytes(&mut out, self.vault_id.as_bytes());
        cbor::text(&mut out, "device_id");
        cbor::bytes(&mut out, self.device_id.as_bytes());
        cbor::text(&mut out, "created_at");
        cbor::uint(&mut out, self.created_at);
        cbor::text(&mut out, "head_revision");
        cbor::uint(&mut out, self.head_revision);
        out
    }

    /// Strict decode: the input must be exactly the canonical encoding of a valid body.
    pub fn decode(bytes: &[u8]) -> Result<Self, ManifestError> {
        use ciborium::Value;
        let v: Value = ciborium::de::from_reader(bytes).map_err(|_| ManifestError::Malformed)?;
        let Value::Map(entries) = v else { return Err(ManifestError::Malformed) };
        let get = |k: &str| entries.iter().find(|(key, _)| key.as_text() == Some(k)).map(|(_, v)| v);
        let b32 = |k: &str| -> Result<[u8; 32], ManifestError> {
            get(k).and_then(Value::as_bytes).and_then(|b| b.as_slice().try_into().ok()).ok_or(ManifestError::Malformed)
        };
        let id = |k: &str| -> Result<Uuid, ManifestError> {
            get(k).and_then(Value::as_bytes).and_then(|b| Uuid::from_slice(b).ok()).ok_or(ManifestError::Malformed)
        };
        let uint = |k: &str| -> Result<u64, ManifestError> {
            get(k).and_then(Value::as_integer).and_then(|i| u64::try_from(i).ok()).ok_or(ManifestError::Malformed)
        };
        let body = ManifestBody {
            vault_id: id("vault_id")?,
            version: uint("version")?,
            head_revision: uint("head_revision")?,
            items: b32("items")?,
            prev: b32("prev")?,
            device_id: id("device_id")?,
            created_at: uint("created_at")?,
        };
        if body.encode() != bytes {
            return Err(ManifestError::Malformed);
        }
        Ok(body)
    }

    pub fn sign(&self, keyset: &AccountKeyset) -> Manifest {
        let body = self.encode();
        let signature = keyset.sign(&body).to_vec();
        Manifest { version: self.version, body, signature }
    }
}

pub fn manifest_hash(body: &[u8]) -> [u8; 32] {
    *blake3::hash(body).as_bytes()
}

/// Signature, shape and vault check. Chain checks are in [`check_chain`].
pub fn verify(sig_pk: &[u8; 32], m: &Manifest, vault_id: Uuid) -> Result<ManifestBody, ManifestError> {
    let sig: [u8; 64] = m.signature.as_slice().try_into().map_err(|_| ManifestError::BadSignature)?;
    verify_signature(sig_pk, &m.body, &sig).map_err(|_| ManifestError::BadSignature)?;
    let body = ManifestBody::decode(&m.body)?;
    if body.vault_id != vault_id {
        return Err(ManifestError::WrongVault);
    }
    if body.version != m.version {
        return Err(ManifestError::Malformed);
    }
    Ok(body)
}

/// Rollback / fork / chain checks against the last manifest this device accepted.
pub fn check_chain(
    seen_version: u64,
    seen_hash: Option<[u8; 32]>,
    body: &ManifestBody,
    hash: [u8; 32],
) -> Result<(), ManifestError> {
    if body.version < seen_version {
        return Err(ManifestError::Rollback { seen: seen_version, got: body.version });
    }
    if body.version == seen_version && seen_version > 0 {
        return if seen_hash == Some(hash) { Ok(()) } else { Err(ManifestError::Fork(body.version)) };
    }
    if body.version == seen_version + 1 {
        let expected = if seen_version == 0 { [0u8; 32] } else { seen_hash.ok_or(ManifestError::ChainBroken)? };
        // A device that never saw version 1 (seen = 0) cannot check the chain of a later manifest,
        // but version 1 itself must start from zeros.
        if body.prev != expected {
            return Err(ManifestError::ChainBroken);
        }
    }
    // version > seen + 1: intermediate manifests were replaced on the server; only the signature
    // and monotonicity can be checked.
    Ok(())
}

#[cfg(test)]
mod tests {
    use kv_proto::{Alg, FORMAT_VERSION};

    use super::*;

    fn body(version: u64, prev: [u8; 32]) -> ManifestBody {
        ManifestBody {
            vault_id: Uuid::from_bytes([1; 16]),
            version,
            head_revision: 9,
            items: [3; 32],
            prev,
            device_id: Uuid::from_bytes([2; 16]),
            created_at: 1_700_000_000,
        }
    }

    #[test]
    fn canonical_roundtrip_and_strictness() {
        let b = body(1, [0; 32]);
        let enc = b.encode();
        assert_eq!(ManifestBody::decode(&enc).unwrap(), b);
        // Same content via a generic (non-canonical key order) encoder is rejected.
        let mut generic = Vec::new();
        ciborium::ser::into_writer(
            &ciborium::Value::Map(vec![
                (ciborium::Value::Text("version".into()), ciborium::Value::Integer(1.into())),
                (ciborium::Value::Text("prev".into()), ciborium::Value::Bytes(vec![0; 32])),
            ]),
            &mut generic,
        )
        .unwrap();
        assert_eq!(ManifestBody::decode(&generic), Err(ManifestError::Malformed));
        let mut trailing = enc.clone();
        trailing.push(0);
        assert!(ManifestBody::decode(&trailing).is_err());
    }

    #[test]
    fn sign_verify() {
        let ks = AccountKeyset::generate(Uuid::new_v4());
        let m = body(1, [0; 32]).sign(&ks);
        assert!(verify(&ks.sig_public(), &m, Uuid::from_bytes([1; 16])).is_ok());
        assert_eq!(verify(&ks.sig_public(), &m, Uuid::nil()), Err(ManifestError::WrongVault));
        let mut forged = m.clone();
        forged.body[5] ^= 1;
        assert_eq!(verify(&ks.sig_public(), &forged, Uuid::from_bytes([1; 16])), Err(ManifestError::BadSignature));
        let other = AccountKeyset::generate(Uuid::new_v4());
        assert_eq!(verify(&other.sig_public(), &m, Uuid::from_bytes([1; 16])), Err(ManifestError::BadSignature));
    }

    #[test]
    fn chain_rules() {
        let b1 = body(1, [0; 32]);
        let h1 = manifest_hash(&b1.encode());
        assert!(check_chain(0, None, &b1, h1).is_ok());
        assert!(check_chain(1, Some(h1), &b1, h1).is_ok());
        let b2 = body(2, h1);
        let h2 = manifest_hash(&b2.encode());
        assert!(check_chain(1, Some(h1), &b2, h2).is_ok());
        assert_eq!(check_chain(2, Some(h2), &b1, h1), Err(ManifestError::Rollback { seen: 2, got: 1 }));
        assert_eq!(check_chain(1, Some(h1), &body(2, [9; 32]), [0; 32]), Err(ManifestError::ChainBroken));
        assert_eq!(check_chain(2, Some(h2), &body(2, h1), [8; 32]), Err(ManifestError::Fork(2)));
        assert!(check_chain(1, Some(h1), &body(5, [7; 32]), [0; 32]).is_ok());
    }

    #[test]
    fn merkle() {
        assert_eq!(merkle_root(vec![]), [0; 32]);
        let a = (Uuid::from_bytes([1; 16]), 1, [1; 32]);
        let b = (Uuid::from_bytes([2; 16]), 2, [2; 32]);
        let c = (Uuid::from_bytes([3; 16]), 3, [3; 32]);
        let r = merkle_root(vec![c, a, b]);
        assert_eq!(r, merkle_root(vec![a, b, c]), "order independent");
        assert_ne!(r, merkle_root(vec![a, b]));
        assert_ne!(r, merkle_root(vec![a, b, (c.0, 4, c.2)]), "revision bound");
        assert_ne!(r, merkle_root(vec![a, b, (c.0, 3, [4; 32])]), "content bound");
    }

    #[test]
    fn envelope_cbor_is_deterministic() {
        let e = Envelope {
            v: FORMAT_VERSION,
            alg: Alg::XChaCha20Poly1305,
            kid: Uuid::nil(),
            n: vec![1; 24],
            ct: vec![2; 3],
            enc: vec![],
        };
        let enc = envelope_cbor(&e);
        assert_eq!(&enc[..3], &[0xa5, 0x61, b'n']);
        let v: ciborium::Value = ciborium::de::from_reader(enc.as_slice()).unwrap();
        assert_eq!(v.as_map().unwrap().len(), 5);
    }
}