//! Key hierarchy (docs/02 "Key hierarchy"):
//! AUK → account keyset → (HPKE) vault key → item key → item payload.

use ed25519_dalek::{Signature, Signer, SigningKey, Verifier, VerifyingKey};
use kv_proto::Envelope;
use rand_core::OsRng;
use uuid::Uuid;
use zeroize::Zeroizing;

use crate::{
    aad::{Aad, ObjectType},
    aead::{pad, unpad, SymKey},
    seal::{seal_to, EncKeypair},
    CryptoError, Result,
};

/// The account's long-term keys. Only ever stored wrapped by the AUK (or a recovery code / LUK).
pub struct AccountKeyset {
    pub account_id: Uuid,
    pub enc: EncKeypair,
    pub sig: SigningKey,
}

impl AccountKeyset {
    pub fn generate(account_id: Uuid) -> Self {
        Self { account_id, enc: EncKeypair::generate(), sig: SigningKey::generate(&mut OsRng) }
    }

    /// Layout: enc key id (16) || X25519 sk (32) || Ed25519 seed (32).
    fn to_bytes(&self) -> Zeroizing<Vec<u8>> {
        let mut out = Zeroizing::new(Vec::with_capacity(80));
        out.extend_from_slice(self.enc.id.as_bytes());
        out.extend_from_slice(&self.enc.secret_bytes());
        out.extend_from_slice(self.sig.as_bytes());
        out
    }

    fn from_bytes(account_id: Uuid, b: &[u8]) -> Result<Self> {
        if b.len() != 80 {
            return Err(CryptoError::Malformed);
        }
        let enc_id = Uuid::from_slice(&b[..16]).map_err(|_| CryptoError::Malformed)?;
        let enc = EncKeypair::from_secret_bytes(enc_id, &b[16..48])?;
        let seed: [u8; 32] = b[48..80].try_into().map_err(|_| CryptoError::Malformed)?;
        Ok(Self { account_id, enc, sig: SigningKey::from_bytes(&seed) })
    }

    /// Wraps with a 32-byte key (AUK, local unlock key or recovery key).
    pub fn wrap(&self, wrapping_key: &SymKey) -> Result<Envelope> {
        wrapping_key.seal(&self.to_bytes(), &Aad::build(ObjectType::Keyset, self.account_id, None, None))
    }

    pub fn unwrap(account_id: Uuid, wrapping_key: &SymKey, env: &Envelope) -> Result<Self> {
        let raw = wrapping_key.open(env, &Aad::build(ObjectType::Keyset, account_id, None, None))?;
        Self::from_bytes(account_id, &raw)
    }

    pub fn sig_public(&self) -> [u8; 32] {
        self.sig.verifying_key().to_bytes()
    }

    pub fn sign(&self, msg: &[u8]) -> [u8; 64] {
        self.sig.sign(msg).to_bytes()
    }

    /// Creates a fresh vault key and seals it to this account.
    pub fn new_vault_key(&self, vault_id: Uuid) -> Result<(SymKey, Envelope)> {
        let vk = SymKey::generate();
        let sealed = self.seal_vault_key(vault_id, &vk)?;
        Ok((vk, sealed))
    }

    pub fn seal_vault_key(&self, vault_id: Uuid, vk: &SymKey) -> Result<Envelope> {
        let mut pt = Zeroizing::new(Vec::with_capacity(48));
        pt.extend_from_slice(vk.id.as_bytes());
        pt.extend_from_slice(vk.expose());
        seal_to(
            self.enc.id,
            &self.enc.public_bytes(),
            &pt,
            &Aad::build(ObjectType::VaultKey, self.account_id, Some(vault_id), None),
        )
    }

    pub fn open_vault_key(&self, vault_id: Uuid, env: &Envelope) -> Result<SymKey> {
        let pt = self.enc.open(env, &Aad::build(ObjectType::VaultKey, self.account_id, Some(vault_id), None))?;
        if pt.len() != 48 {
            return Err(CryptoError::Malformed);
        }
        let id = Uuid::from_slice(&pt[..16]).map_err(|_| CryptoError::Malformed)?;
        let key: [u8; 32] = pt[16..].try_into().map_err(|_| CryptoError::Malformed)?;
        Ok(SymKey::from_bytes(id, key))
    }
}

pub fn verify_signature(public: &[u8; 32], msg: &[u8], sig: &[u8; 64]) -> Result<()> {
    let vk = VerifyingKey::from_bytes(public).map_err(|_| CryptoError::Malformed)?;
    vk.verify(msg, &Signature::from_bytes(sig)).map_err(|_| CryptoError::Decrypt)
}

/// An encrypted item as stored locally and on the server.
pub struct SealedItem {
    pub wrapped_item_key: Envelope,
    pub payload: Envelope,
}

/// Encrypts an item's plaintext (already serialized) with a fresh per-item key.
pub fn seal_item(
    account_id: Uuid,
    vault_id: Uuid,
    item_id: Uuid,
    vault_key: &SymKey,
    plaintext: &[u8],
) -> Result<SealedItem> {
    let item_key = SymKey::generate();
    let mut key_pt = Zeroizing::new(Vec::with_capacity(48));
    key_pt.extend_from_slice(item_key.id.as_bytes());
    key_pt.extend_from_slice(item_key.expose());
    let wrapped_item_key =
        vault_key.seal(&key_pt, &Aad::build(ObjectType::ItemKey, account_id, Some(vault_id), Some(item_id)))?;
    let payload =
        item_key.seal(&pad(plaintext), &Aad::build(ObjectType::Item, account_id, Some(vault_id), Some(item_id)))?;
    Ok(SealedItem { wrapped_item_key, payload })
}

pub fn open_item(
    account_id: Uuid,
    vault_id: Uuid,
    item_id: Uuid,
    vault_key: &SymKey,
    sealed: &SealedItem,
) -> Result<Zeroizing<Vec<u8>>> {
    let key_pt = vault_key
        .open(&sealed.wrapped_item_key, &Aad::build(ObjectType::ItemKey, account_id, Some(vault_id), Some(item_id)))?;
    if key_pt.len() != 48 {
        return Err(CryptoError::Malformed);
    }
    let id = Uuid::from_slice(&key_pt[..16]).map_err(|_| CryptoError::Malformed)?;
    let item_key = SymKey::from_bytes(id, key_pt[16..].try_into().map_err(|_| CryptoError::Malformed)?);
    let padded =
        item_key.open(&sealed.payload, &Aad::build(ObjectType::Item, account_id, Some(vault_id), Some(item_id)))?;
    unpad(&padded)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{kdf, secret_key::SecretKey};

    #[test]
    fn full_hierarchy() {
        let account = Uuid::new_v4();
        let sk = SecretKey::generate();
        let derived = kdf::derive_keys_unchecked("hunter2", &sk, account, &kdf::test_params()).unwrap();
        let auk = SymKey::from_bytes(Uuid::nil(), *derived.auk);

        let keyset = AccountKeyset::generate(account);
        let wrapped = keyset.wrap(&auk).unwrap();
        let keyset = AccountKeyset::unwrap(account, &auk, &wrapped).unwrap();

        let vault = Uuid::new_v4();
        let (_, sealed_vk) = keyset.new_vault_key(vault).unwrap();
        let vk = keyset.open_vault_key(vault, &sealed_vk).unwrap();

        let item = Uuid::new_v4();
        let sealed = seal_item(account, vault, item, &vk, br#"{"title":"bank"}"#).unwrap();
        assert_eq!(&**open_item(account, vault, item, &vk, &sealed).unwrap(), br#"{"title":"bank"}"#);

        // The server swaps the ciphertext into another item id, which is detected.
        assert!(open_item(account, vault, Uuid::new_v4(), &vk, &sealed).is_err());
        // A vault key sealed for one vault cannot be opened as another vault's key.
        assert!(keyset.open_vault_key(Uuid::new_v4(), &sealed_vk).is_err());
    }

    #[test]
    fn wrong_password_cannot_unwrap() {
        let account = Uuid::new_v4();
        let sk = SecretKey::generate();
        let good = kdf::derive_keys_unchecked("right", &sk, account, &kdf::test_params()).unwrap();
        let bad = kdf::derive_keys_unchecked("wrong", &sk, account, &kdf::test_params()).unwrap();
        let wrapped = AccountKeyset::generate(account).wrap(&SymKey::from_bytes(Uuid::nil(), *good.auk)).unwrap();
        assert!(AccountKeyset::unwrap(account, &SymKey::from_bytes(Uuid::nil(), *bad.auk), &wrapped).is_err());
    }

    #[test]
    fn signatures() {
        let ks = AccountKeyset::generate(Uuid::new_v4());
        let sig = ks.sign(b"manifest");
        assert!(verify_signature(&ks.sig_public(), b"manifest", &sig).is_ok());
        assert!(verify_signature(&ks.sig_public(), b"manifesT", &sig).is_err());
    }
}