//! The invitation inbox — F5's exact shapes: recipient-addressed sealed
//! envelopes in a bounded multi-writer contract with a deterministic merge
//! (validate → dedup → caps → tombstones), an owner PruneRecord, and the
//! owner-signed InboxDescriptor that says how to reach a persona. BYTE-
//! COMPATIBILITY IS THE CONTRACT: the bundled inbox / descriptor wasm
//! validate exactly these (built from research/freenet-f5, code hashes
//! pinned by the app's tests).
//!
//! Signing rides `soulstream.friend/v1` — the delegate generation in the
//! app has no inbox domain yet (F5's deferred item); the contracts were
//! built to that. Sealing: X25519 (sender ephemeral → recipient box key)
//! into ChaCha20-Poly1305 with the signed header as AAD.

use ed25519_dalek::VerifyingKey;
use serde::{Deserialize, Serialize};
use std::collections::BTreeMap;

pub const SCHEMA_VERSION: u32 = 1;
pub const INBOX_PARAMS_FAMILY: &[u8] = b"soulstream.inbox.v1";
pub const DESC_PARAMS_FAMILY: &[u8] = b"soulstream.inbox-desc.v1";
pub const ENVELOPE_DOMAIN: &str = "soulstream.friend/v1";
pub const DESC_DOMAIN: &str = "soulstream.friend/v1";

pub const MAX_STATE_BYTES: usize = 1024 * 1024;
pub const MAX_ENVELOPES: usize = 64;
pub const PER_SENDER_CAP: usize = 4;
pub const ANON_POOL_CAP: usize = 8;
pub const MAX_CIPHERTEXT: usize = 8 * 1024;
pub const MAX_AAD: usize = 256;
pub const MAX_LIFETIME_SECS: u64 = 30 * 24 * 3600;
pub const MAX_TOMBSTONES: usize = 256;
pub const ANON_POW_BITS: u32 = 12;

// ---------------------------------------------------------------- envelope

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum SenderClass {
    Known([u8; 32]),
    Anonymous,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Envelope {
    pub envelope_id: [u8; 16],
    pub recipient: [u8; 32],
    pub sender: SenderClass,
    pub created_at: u64,
    pub expires_at: u64,
    /// Known: empty. Anonymous: a nonce whose blake3(id ‖ nonce) carries
    /// ANON_POW_BITS leading zero bits.
    pub admission: Vec<u8>,
    /// Sender ephemeral X25519 public key (32) ‖ ChaCha20-Poly1305
    /// ciphertext of the cargo.
    pub sealed: Vec<u8>,
    /// Known sender's signature over the header (delegate-signed).
    pub sender_sig: Vec<u8>,
}

#[derive(Serialize)]
struct EnvelopeHeader<'a> {
    envelope_id: &'a [u8; 16],
    recipient: &'a [u8; 32],
    sender: &'a SenderClass,
    created_at: u64,
    expires_at: u64,
}

/// The signed / authenticated header — also the AAD for the seal.
pub fn envelope_header_bytes(e: &Envelope) -> Vec<u8> {
    bincode::serialize(&EnvelopeHeader {
        envelope_id: &e.envelope_id,
        recipient: &e.recipient,
        sender: &e.sender,
        created_at: e.created_at,
        expires_at: e.expires_at,
    })
    .expect("serializes")
}

/// What the delegate is handed to sign for a header / prune / descriptor.
pub fn sign_payload(body: &[u8]) -> Vec<u8> {
    blake3::hash(body).as_bytes().to_vec()
}

pub fn box_key(shared_secret: &[u8; 32]) -> [u8; 32] {
    blake3::derive_key("soulstream.inbox-box.v1", shared_secret)
}
pub fn box_nonce(eph_pub: &[u8; 32], envelope_id: &[u8; 16]) -> [u8; 12] {
    let mut input = Vec::with_capacity(48);
    input.extend_from_slice(eph_pub);
    input.extend_from_slice(envelope_id);
    let h = blake3::hash(&input);
    h.as_bytes()[..12].try_into().unwrap()
}

pub fn anon_admission_ok(envelope_id: &[u8; 16], admission: &[u8]) -> bool {
    if admission.len() != 16 {
        return false;
    }
    let mut input = Vec::with_capacity(32);
    input.extend_from_slice(envelope_id);
    input.extend_from_slice(admission);
    let h = blake3::hash(&input);
    let mut zeros = 0u32;
    for b in h.as_bytes() {
        if *b == 0 {
            zeros += 8;
        } else {
            zeros += b.leading_zeros();
            break;
        }
    }
    zeros >= ANON_POW_BITS
}

#[derive(Debug, PartialEq, Eq)]
pub enum EnvelopeInvalid {
    AadTooBig(usize),
    CiphertextTooBig(usize),
    WrongRecipient,
    LifetimeInsane,
    BadSenderKey,
    BadSenderSig,
    BadAdmission,
    SealedTooShort,
}

/// Structural + authorization validation, cheapest first.
pub fn validate_envelope(owner: &[u8; 32], e: &Envelope) -> Result<(), EnvelopeInvalid> {
    if e.sealed.len() > 32 + MAX_CIPHERTEXT {
        return Err(EnvelopeInvalid::CiphertextTooBig(e.sealed.len()));
    }
    if e.sealed.len() < 32 + 16 {
        return Err(EnvelopeInvalid::SealedTooShort);
    }
    if e.admission.len() > MAX_AAD {
        return Err(EnvelopeInvalid::AadTooBig(e.admission.len()));
    }
    if &e.recipient != owner {
        return Err(EnvelopeInvalid::WrongRecipient);
    }
    if e.expires_at <= e.created_at || e.expires_at - e.created_at > MAX_LIFETIME_SECS {
        return Err(EnvelopeInvalid::LifetimeInsane);
    }
    match &e.sender {
        SenderClass::Known(key) => {
            if VerifyingKey::from_bytes(key).is_err() {
                return Err(EnvelopeInvalid::BadSenderKey);
            }
            if !crate::verify_sig_in(
                ENVELOPE_DOMAIN,
                key,
                &sign_payload(&envelope_header_bytes(e)),
                &e.sender_sig,
            ) {
                return Err(EnvelopeInvalid::BadSenderSig);
            }
        }
        SenderClass::Anonymous => {
            if !anon_admission_ok(&e.envelope_id, &e.admission) {
                return Err(EnvelopeInvalid::BadAdmission);
            }
        }
    }
    Ok(())
}

// ---------------------------------------------------------------- sealing

/// Seal `cargo` to a recipient's box key. `eph_secret` is 32 fresh random
/// bytes from the caller (this crate draws no entropy of its own).
pub fn seal(
    recipient_box_pub: &[u8; 32],
    envelope_id: &[u8; 16],
    header_aad: &[u8],
    cargo: &[u8],
    eph_secret: [u8; 32],
) -> Vec<u8> {
    use chacha20poly1305::aead::{Aead, KeyInit, Payload};
    use chacha20poly1305::{ChaCha20Poly1305, Key, Nonce};
    let eph = x25519_dalek::StaticSecret::from(eph_secret);
    let eph_pub = x25519_dalek::PublicKey::from(&eph);
    let shared = eph.diffie_hellman(&x25519_dalek::PublicKey::from(*recipient_box_pub));
    let key = box_key(shared.as_bytes());
    let nonce = box_nonce(eph_pub.as_bytes(), envelope_id);
    let ct = ChaCha20Poly1305::new(Key::from_slice(&key))
        .encrypt(Nonce::from_slice(&nonce), Payload { msg: cargo, aad: header_aad })
        .expect("seal");
    let mut out = eph_pub.as_bytes().to_vec();
    out.extend_from_slice(&ct);
    out
}

/// Open an envelope with the recipient's box secret. None = not for this
/// key, or tampered — the caller quarantines, never crashes.
pub fn open(box_secret: &[u8; 32], e: &Envelope) -> Option<Vec<u8>> {
    use chacha20poly1305::aead::{Aead, KeyInit, Payload};
    use chacha20poly1305::{ChaCha20Poly1305, Key, Nonce};
    if e.sealed.len() < 32 {
        return None;
    }
    let eph_pub: [u8; 32] = e.sealed[..32].try_into().ok()?;
    let secret = x25519_dalek::StaticSecret::from(*box_secret);
    let shared = secret.diffie_hellman(&x25519_dalek::PublicKey::from(eph_pub));
    let key = box_key(shared.as_bytes());
    let nonce = box_nonce(&eph_pub, &e.envelope_id);
    ChaCha20Poly1305::new(Key::from_slice(&key))
        .decrypt(
            Nonce::from_slice(&nonce),
            Payload { msg: &e.sealed[32..], aad: &envelope_header_bytes(e) },
        )
        .ok()
}

/// The X25519 public key for a box secret.
pub fn box_public(box_secret: &[u8; 32]) -> [u8; 32] {
    let s = x25519_dalek::StaticSecret::from(*box_secret);
    *x25519_dalek::PublicKey::from(&s).as_bytes()
}

// ---------------------------------------------------------------- state

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PruneRecord {
    pub seq: u64,
    pub removed_ids: Vec<[u8; 16]>,
    /// Everything with expires_at <= this floor stays out permanently.
    pub floor_expiry: u64,
    pub sig: Vec<u8>,
}

#[derive(Serialize)]
struct PruneSigBody<'a> {
    seq: u64,
    removed_ids: &'a [[u8; 16]],
    floor_expiry: u64,
}

pub fn prune_sig_body(p: &PruneRecord) -> Vec<u8> {
    bincode::serialize(&PruneSigBody {
        seq: p.seq,
        removed_ids: &p.removed_ids,
        floor_expiry: p.floor_expiry,
    })
    .expect("serializes")
}

pub fn prune_valid(owner: &[u8; 32], p: &PruneRecord) -> bool {
    p.removed_ids.len() <= MAX_TOMBSTONES
        && crate::verify_sig_in(ENVELOPE_DOMAIN, owner, &sign_payload(&prune_sig_body(p)), &p.sig)
}

#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
pub struct InboxState {
    pub schema_version: u32,
    pub owner: [u8; 32],
    pub envelopes: Vec<Envelope>,
    pub prune: Option<PruneRecord>,
}

pub fn make_inbox_params(owner: &[u8; 32]) -> Vec<u8> {
    let mut p = Vec::with_capacity(INBOX_PARAMS_FAMILY.len() + 32);
    p.extend_from_slice(INBOX_PARAMS_FAMILY);
    p.extend_from_slice(owner);
    p
}

/// Deterministic multi-writer merge. Inputs are RAW candidate states (any
/// mix of replicas and single-envelope offers); output is the bounded
/// canonical state. Mirrors the bundled contract exactly.
pub fn merge_inbox(owner: &[u8; 32], candidates: &[&[u8]]) -> InboxState {
    let mut best_prune: Option<(u64, [u8; 32], PruneRecord)> = None;
    let mut parsed: Vec<InboxState> = Vec::new();
    for bytes in candidates {
        if bytes.is_empty() || bytes.len() > MAX_STATE_BYTES {
            continue;
        }
        let Ok(st) = bincode::deserialize::<InboxState>(bytes) else { continue };
        if st.schema_version != SCHEMA_VERSION || &st.owner != owner {
            continue;
        }
        if let Some(p) = &st.prune {
            if prune_valid(owner, p) {
                let rank = (p.seq, *blake3::hash(&prune_sig_body(p)).as_bytes());
                let better = match &best_prune {
                    None => true,
                    Some((s, h, _)) => rank > (*s, *h),
                };
                if better {
                    best_prune = Some((rank.0, rank.1, p.clone()));
                }
            }
        }
        parsed.push(st);
    }
    let prune = best_prune.map(|(_, _, p)| p);

    let mut by_id: BTreeMap<[u8; 16], (Envelope, [u8; 32])> = BTreeMap::new();
    for st in &parsed {
        for e in &st.envelopes {
            if validate_envelope(owner, e).is_err() {
                continue;
            }
            if let Some(p) = &prune {
                if p.removed_ids.contains(&e.envelope_id) || e.expires_at <= p.floor_expiry {
                    continue;
                }
            }
            let h = *blake3::hash(&bincode::serialize(e).unwrap()).as_bytes();
            match by_id.get(&e.envelope_id) {
                Some((_, existing_h)) if existing_h <= &h => {}
                _ => {
                    by_id.insert(e.envelope_id, (e.clone(), h));
                }
            }
        }
    }

    let mut all: Vec<Envelope> = by_id.into_values().map(|(e, _)| e).collect();
    all.sort_by_key(|e| (e.created_at, e.envelope_id));
    let mut per_sender: BTreeMap<[u8; 32], Vec<Envelope>> = BTreeMap::new();
    let mut anon_pool: Vec<Envelope> = Vec::new();
    for e in all {
        match &e.sender {
            SenderClass::Known(k) => per_sender.entry(*k).or_default().push(e),
            SenderClass::Anonymous => {
                if anon_pool.len() < ANON_POOL_CAP {
                    anon_pool.push(e);
                }
            }
        }
    }
    let mut survivors: Vec<Envelope> = anon_pool;
    for (_, mut list) in per_sender {
        let keep_from = list.len().saturating_sub(PER_SENDER_CAP);
        survivors.extend(list.drain(keep_from..));
    }
    survivors.sort_by_key(|e| (e.created_at, e.envelope_id));
    survivors.truncate(MAX_ENVELOPES);
    InboxState {
        schema_version: SCHEMA_VERSION,
        owner: *owner,
        envelopes: survivors,
        prune,
    }
}

// ---------------------------------------------------------------- descriptor

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct InboxDescriptor {
    pub schema_version: u32,
    pub owner: [u8; 32],
    pub sequence: u64,
    pub inbox_family: Vec<u8>,
    pub box_public_key: [u8; 32],
    pub box_key_id: u32,
    pub max_envelope_bytes: u32,
    pub admission_policy: Vec<String>,
    pub sig: Vec<u8>,
}

#[derive(Serialize)]
struct DescSigBody<'a> {
    schema_version: u32,
    owner: &'a [u8; 32],
    sequence: u64,
    inbox_family: &'a [u8],
    box_public_key: &'a [u8; 32],
    box_key_id: u32,
    max_envelope_bytes: u32,
    admission_policy: &'a [String],
}

pub fn desc_sig_body(d: &InboxDescriptor) -> Vec<u8> {
    bincode::serialize(&DescSigBody {
        schema_version: d.schema_version,
        owner: &d.owner,
        sequence: d.sequence,
        inbox_family: &d.inbox_family,
        box_public_key: &d.box_public_key,
        box_key_id: d.box_key_id,
        max_envelope_bytes: d.max_envelope_bytes,
        admission_policy: &d.admission_policy,
    })
    .expect("serializes")
}

pub fn make_desc_params(owner: &[u8; 32]) -> Vec<u8> {
    let mut p = Vec::with_capacity(DESC_PARAMS_FAMILY.len() + 32);
    p.extend_from_slice(DESC_PARAMS_FAMILY);
    p.extend_from_slice(owner);
    p
}

pub fn validate_descriptor(owner: &[u8; 32], bytes: &[u8]) -> Option<InboxDescriptor> {
    if bytes.len() > 4096 {
        return None;
    }
    let d: InboxDescriptor = bincode::deserialize(bytes).ok()?;
    if d.schema_version != SCHEMA_VERSION || &d.owner != owner {
        return None;
    }
    crate::verify_sig_in(DESC_DOMAIN, owner, &sign_payload(&desc_sig_body(&d)), &d.sig).then_some(d)
}

#[cfg(test)]
mod tests {
    use super::*;
    use ed25519_dalek::{Signer, SigningKey};

    fn sign(key: &SigningKey, payload: &[u8]) -> Vec<u8> {
        key.sign(&crate::signable(ENVELOPE_DOMAIN, payload)).to_bytes().to_vec()
    }

    #[test]
    fn seal_open_merge_and_prune_hold_their_gate() {
        let recipient = SigningKey::from_bytes(&[1; 32]);
        let sender = SigningKey::from_bytes(&[2; 32]);
        let owner = recipient.verifying_key().to_bytes();
        let box_secret = [7u8; 32];
        let box_pub = box_public(&box_secret);

        // a sealed, signed envelope opens with the right key and not another
        let mut e = Envelope {
            envelope_id: [3; 16],
            recipient: owner,
            sender: SenderClass::Known(sender.verifying_key().to_bytes()),
            created_at: 100,
            expires_at: 100 + 3600,
            admission: vec![],
            sealed: vec![],
            sender_sig: vec![],
        };
        e.sealed = seal(&box_pub, &e.envelope_id, &envelope_header_bytes(&e), b"hello", [9u8; 32]);
        e.sender_sig = sign(&sender, &sign_payload(&envelope_header_bytes(&e)));
        assert_eq!(validate_envelope(&owner, &e), Ok(()));
        assert_eq!(open(&box_secret, &e).as_deref(), Some(&b"hello"[..]));
        assert_eq!(open(&[8u8; 32], &e), None);
        // a tampered header breaks both the signature and the AAD
        let mut t = e.clone();
        t.expires_at += 1;
        assert_eq!(validate_envelope(&owner, &t), Err(EnvelopeInvalid::BadSenderSig));
        assert_eq!(open(&box_secret, &t), None);

        // merge: an offer state with the envelope lands; a prune removes it for good
        let offer = bincode::serialize(&InboxState { schema_version: SCHEMA_VERSION, owner, envelopes: vec![e.clone()], prune: None }).unwrap();
        let merged = merge_inbox(&owner, &[&offer]);
        assert_eq!(merged.envelopes.len(), 1);
        let mut prune = PruneRecord { seq: 1, removed_ids: vec![e.envelope_id], floor_expiry: 0, sig: vec![] };
        prune.sig = sign(&recipient, &sign_payload(&prune_sig_body(&prune)));
        let pruned = bincode::serialize(&InboxState { schema_version: SCHEMA_VERSION, owner, envelopes: vec![], prune: Some(prune) }).unwrap();
        let after = merge_inbox(&owner, &[&offer, &pruned]);
        assert!(after.envelopes.is_empty());
        assert_eq!(after.prune.as_ref().map(|p| p.seq), Some(1));
        // a forged prune has no authority
        let mut forged = PruneRecord { seq: 5, removed_ids: vec![e.envelope_id], floor_expiry: 0, sig: vec![] };
        forged.sig = sign(&sender, &sign_payload(&prune_sig_body(&forged)));
        let fb = bincode::serialize(&InboxState { schema_version: SCHEMA_VERSION, owner, envelopes: vec![], prune: Some(forged) }).unwrap();
        assert_eq!(merge_inbox(&owner, &[&offer, &fb]).envelopes.len(), 1);

        // descriptor round trip
        let mut d = InboxDescriptor { schema_version: SCHEMA_VERSION, owner, sequence: 1, inbox_family: INBOX_PARAMS_FAMILY.to_vec(), box_public_key: box_pub, box_key_id: 1, max_envelope_bytes: MAX_CIPHERTEXT as u32, admission_policy: vec!["known-signed".into()], sig: vec![] };
        d.sig = sign(&recipient, &sign_payload(&desc_sig_body(&d)));
        let db = bincode::serialize(&d).unwrap();
        assert_eq!(validate_descriptor(&owner, &db), Some(d));
    }
}