//! Crypto-shred end-to-end: write encrypted bytes; destroy key; decrypt fails.

use klieo_ops::crypter::{
    Crypter, CrypterError, KeyProvider, StaticKeyProvider, TenantKeyedCrypter,
};
use klieo_ops::tenant::scope_tenant;
use klieo_ops::types::TenantId;
use std::collections::HashMap;
use std::sync::{Arc, RwLock};

/// [`KeyProvider`] whose keys can be removed mid-test to simulate crypto-shred.
struct MutableKeyProvider {
    inner: RwLock<HashMap<Option<TenantId>, [u8; 32]>>,
}

impl MutableKeyProvider {
    fn new(keys: HashMap<Option<TenantId>, [u8; 32]>) -> Self {
        Self {
            inner: RwLock::new(keys),
        }
    }

    fn shred(&self, tenant: &Option<TenantId>) {
        self.inner.write().expect("lock poisoned").remove(tenant);
    }
}

impl KeyProvider for MutableKeyProvider {
    fn lookup(&self, tenant: &Option<TenantId>) -> Option<[u8; 32]> {
        self.inner
            .read()
            .expect("lock poisoned")
            .get(tenant)
            .copied()
    }
}

#[tokio::test]
async fn shred_key_makes_past_ciphertext_unreadable() {
    let mut map = HashMap::new();
    map.insert(Some(TenantId("BL_auto".into())), [42u8; 32]);
    let keys = Arc::new(MutableKeyProvider::new(map));
    let crypter = TenantKeyedCrypter::new(keys.clone());

    let plaintext = b"customer-PII-payload";
    let ciphertext = scope_tenant(Some(TenantId("BL_auto".into())), async {
        crypter.encrypt(plaintext).await.expect("encrypt")
    })
    .await;

    // Decrypt while key still present — must succeed.
    let recovered = scope_tenant(Some(TenantId("BL_auto".into())), async {
        crypter.decrypt(&ciphertext).await.expect("decrypt")
    })
    .await;
    assert_eq!(recovered, plaintext);

    // Crypto-shred: destroy the tenant's key.
    keys.shred(&Some(TenantId("BL_auto".into())));

    // Decrypt now fails with KeyDestroyed.
    let err = scope_tenant(Some(TenantId("BL_auto".into())), async {
        crypter.decrypt(&ciphertext).await.unwrap_err()
    })
    .await;
    assert!(
        matches!(err, CrypterError::KeyDestroyed { .. }),
        "expected KeyDestroyed, got {err:?}"
    );
}

#[tokio::test]
async fn different_tenants_cannot_decrypt_each_others_ciphertext() {
    let mut map = HashMap::new();
    map.insert(Some(TenantId("alice".into())), [1u8; 32]);
    map.insert(Some(TenantId("bob".into())), [2u8; 32]);
    let crypter = TenantKeyedCrypter::new(Arc::new(StaticKeyProvider::from_map(map)));

    let alice_ct = scope_tenant(Some(TenantId("alice".into())), async {
        crypter.encrypt(b"alice secret").await.expect("encrypt")
    })
    .await;

    // Bob's tenant context does not match Alice's AAD → AadMismatch (detected
    // before the AEAD tag check, which would also fail). AadMismatch is the
    // more precise signal added by the B1 AAD-binding fix.
    let err = scope_tenant(Some(TenantId("bob".into())), async {
        crypter.decrypt(&alice_ct).await.unwrap_err()
    })
    .await;
    assert!(
        matches!(err, CrypterError::AadMismatch),
        "expected AadMismatch for cross-tenant decrypt, got {err:?}"
    );
}
