use anyhow::Result;
use chacha20poly1305::{
    aead::{Aead, KeyInit},
    ChaCha20Poly1305, Key, Nonce,
};
use rand::rngs::OsRng;
use x25519_dalek::{PublicKey, StaticSecret};

pub struct Cipher {
    inner: ChaCha20Poly1305,
}

impl Cipher {
    pub fn new(key_bytes: &[u8; 32]) -> Self {
        let key = Key::from_slice(key_bytes);
        Self {
            inner: ChaCha20Poly1305::new(key),
        }
    }

    pub fn encrypt(&self, nonce: &[u8; 12], plaintext: &[u8]) -> Result<Vec<u8>> {
        let nonce = Nonce::from_slice(nonce);
        self.inner
            .encrypt(nonce, plaintext)
            .map_err(|_e| anyhow::anyhow!("Encryption error"))
    }

    pub fn decrypt(&self, nonce: &[u8; 12], ciphertext: &[u8]) -> Result<Vec<u8>> {
        let nonce = Nonce::from_slice(nonce);
        self.inner
            .decrypt(nonce, ciphertext)
            .map_err(|_e| anyhow::anyhow!("Decryption error"))
    }
}

// X25519 Key Exchange for Session Encryption
pub struct KeyExchange {
    secret: StaticSecret,
}

impl KeyExchange {
    pub fn new() -> Self {
        Self {
            secret: StaticSecret::random_from_rng(OsRng),
        }
    }

    pub fn from_secret_bytes(bytes: [u8; 32]) -> Self {
        Self {
            secret: StaticSecret::from(bytes),
        }
    }

    pub fn public_key_bytes(&self) -> [u8; 32] {
        let pub_key = PublicKey::from(&self.secret);
        *pub_key.as_bytes()
    }

    pub fn exchange(&self, peer_public: &[u8; 32]) -> [u8; 32] {
        let peer_pub = PublicKey::from(*peer_public);
        let shared = self.secret.diffie_hellman(&peer_pub);
        *shared.as_bytes()
    }

    pub fn as_bytes(&self) -> [u8; 32] {
        self.secret.to_bytes()
    }
}

// Ed25519 Identity Keys for Micro-CA Signatures
use ed25519_dalek::{Signer, Verifier};

pub struct IdentityKey {
    keypair: ed25519_dalek::SigningKey,
}

impl IdentityKey {
    pub fn generate() -> Self {
        use rand::RngCore;
        let mut bytes = [0u8; 32];
        OsRng.fill_bytes(&mut bytes);
        Self {
            keypair: ed25519_dalek::SigningKey::from_bytes(&bytes),
        }
    }

    pub fn from_bytes(bytes: &[u8; 32]) -> Self {
        Self {
            keypair: ed25519_dalek::SigningKey::from_bytes(bytes),
        }
    }

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

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

    pub fn as_bytes(&self) -> [u8; 32] {
        self.keypair.to_bytes()
    }
}

pub struct IdentityVerifier;

impl IdentityVerifier {
    pub fn verify(pub_key: &[u8; 32], message: &[u8], signature: &[u8; 64]) -> Result<()> {
        let key = ed25519_dalek::VerifyingKey::from_bytes(pub_key)
            .map_err(|_| anyhow::anyhow!("Invalid Ed25519 public key"))?;
        let sig = ed25519_dalek::Signature::from_bytes(signature);
        key.verify(message, &sig)
            .map_err(|_| anyhow::anyhow!("Signature verification failed"))
    }
}

// ── Key Derivation (Brain Wallet) ────────────────────────────────────────────
//
// Algorithm: Argon2id (memory-hard, side-channel resistant)
// All improvements over the original PBKDF2 path:
//   1. Memory-hard KDF (GPU/ASIC resistance)
//   2. Salt is 32 cryptographically random bytes (not a string)
//   3. Passphrase is NFKC-normalised + trimmed before hashing
//   4. KDF parameters are stored alongside derived data (recoverability)
//   5. Format version field enables future algorithm agility
//   6. Returns Result<> — no panics in crypto code

use argon2::{Algorithm, Argon2, Params, Version};
use rand::RngCore;
use serde::{Deserialize, Serialize};
use unicode_normalization::UnicodeNormalization;

/// Versioned KDF parameters that must be stored alongside any data protected
/// by a key derived from a mantra. They are everything needed to reproduce the
/// exact same key from the same mantra on any platform, now or in the future.
#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct KdfParams {
    /// Format version (currently 2). Increment when changing algorithm or
    /// parameter semantics so old configs can still be read.
    pub version: u8,
    /// KDF algorithm name — always "argon2id" for version 2.
    pub algorithm: String,
    /// Memory cost in KiB (262144 = 256 MiB).
    pub m_cost: u32,
    /// Number of iterations / time cost.
    pub t_cost: u32,
    /// Degree of parallelism.
    pub p_cost: u32,
    /// 32 cryptographically-random bytes, hex-encoded. Never derived from user
    /// input. Must be stored; losing it makes the key irrecoverable.
    pub salt_hex: String,
}

impl KdfParams {
    /// Default secure parameters: RFC 9106 "interactive" profile.
    /// ~256 MiB RAM, 3 passes, 4 threads — takes ≈1–3 s on typical hardware.
    pub const DEFAULT_M_COST: u32 = 262_144; // 256 MiB
    pub const DEFAULT_T_COST: u32 = 3;
    pub const DEFAULT_P_COST: u32 = 4;

    /// Create new v2 params with the given pre-generated random salt.
    pub fn new_v2(salt: [u8; 32]) -> Self {
        Self {
            version: 2,
            algorithm: "argon2id".to_string(),
            m_cost: Self::DEFAULT_M_COST,
            t_cost: Self::DEFAULT_T_COST,
            p_cost: Self::DEFAULT_P_COST,
            salt_hex: hex::encode(salt),
        }
    }
}

pub struct KeyGenerator;

impl KeyGenerator {
    /// Generate a fresh [`KdfParams`] with a cryptographically-random 32-byte
    /// salt. Call this once per bootstrap; then store the params in `server.toml`.
    pub fn generate_kdf_params() -> KdfParams {
        let mut salt = [0u8; 32];
        OsRng.fill_bytes(&mut salt);
        KdfParams::new_v2(salt)
    }

    /// Derive a 32-byte master key from a human passphrase (mantra).
    ///
    /// ## What this does
    /// 1. **NFKC-normalise** the passphrase and trim surrounding whitespace so
    ///    that different Unicode representations of the same text always hash
    ///    identically (critical for cross-platform recovery).
    /// 2. Decode the random salt from `params.salt_hex`.
    /// 3. Run **Argon2id** with the stored memory/time/parallelism parameters.
    ///
    /// Returns `Err` on invalid params or hex — never panics.
    pub fn derive_from_passphrase(passphrase: &str, params: &KdfParams) -> Result<[u8; 32]> {
        // Step 1: NFKC normalise + trim (cross-platform / cross-keyboard safety)
        let normalised: String = passphrase.nfkc().collect::<String>().trim().to_string();
        if normalised.is_empty() {
            anyhow::bail!("Passphrase must not be empty after normalisation");
        }

        // Step 2: Decode the stored random salt
        let salt_bytes = hex::decode(&params.salt_hex)
            .map_err(|e| anyhow::anyhow!("Invalid salt hex: {}", e))?;

        // Step 3: Build Argon2id with stored parameters
        let argon2_params = Params::new(params.m_cost, params.t_cost, params.p_cost, Some(32))
            .map_err(|e| anyhow::anyhow!("Invalid Argon2 parameters: {}", e))?;

        let argon2 = Argon2::new(Algorithm::Argon2id, Version::V0x13, argon2_params);

        let mut key = [0u8; 32];
        argon2
            .hash_password_into(normalised.as_bytes(), &salt_bytes, &mut key)
            .map_err(|e| anyhow::anyhow!("Argon2id derivation failed: {}", e))?;

        Ok(key)
    }

    /// **Legacy v1 path — signature preserved for migration tooling.**
    ///
    /// The pbkdf2/hmac/sha2 crates have been removed. To re-enable this for a
    /// migration binary, add those deps back and replace this body.
    #[allow(dead_code)]
    pub fn derive_from_passphrase_v1_legacy(_passphrase: &str, _salt: &str) -> ! {
        panic!(
            "Legacy PBKDF2 path is not compiled in. \
             Add pbkdf2/hmac/sha2 deps and restore the body to use this."
        )
    }
}

pub struct SessionKeys {
    pub server_to_client: [u8; 32],
    pub client_to_server: [u8; 32],
}

impl SessionKeys {
    pub fn derive(shared_secret: &[u8; 32], client_pub: &[u8; 32], server_pub: &[u8; 32]) -> Self {
        // Simple HKDF-like derivation using BLAKE3
        let mut context = Vec::new();
        context.extend_from_slice(client_pub);
        context.extend_from_slice(server_pub);

        let mut hasher = blake3::Hasher::new_derive_key("konduit-s2c");
        hasher.update(shared_secret);
        hasher.update(&context);
        let mut s2c = [0u8; 32];
        s2c.copy_from_slice(hasher.finalize().as_bytes());

        let mut hasher = blake3::Hasher::new_derive_key("konduit-c2s");
        hasher.update(shared_secret);
        hasher.update(&context);
        let mut c2s = [0u8; 32];
        c2s.copy_from_slice(hasher.finalize().as_bytes());

        SessionKeys {
            server_to_client: s2c,
            client_to_server: c2s,
        }
    }
}

pub struct NonceCounter {
    counter: u64,
}

impl NonceCounter {
    pub fn new() -> Self {
        Self { counter: 0 }
    }

    pub fn next(&mut self) -> [u8; 12] {
        let mut nonce = [0u8; 12];
        // 4 bytes random prefix (salt) to prevent collision on restart?
        // For now just 12 bytes counter le
        nonce[0..8].copy_from_slice(&self.counter.to_le_bytes());
        self.counter += 1;
        nonce
    }
}

/// Compute HMAC-like MAC using BLAKE3 for authentication
pub fn compute_mac(key: &[u8], message: &[u8]) -> Result<[u8; 32]> {
    use blake3::Hasher;

    let key_32: [u8; 32] = key
        .try_into()
        .map_err(|_| anyhow::anyhow!("Mac key must be exactly 32 bytes (64 hex characters)"))?;

    let mut hasher = Hasher::new_keyed(&key_32);
    hasher.update(message);
    Ok(*hasher.finalize().as_bytes())
}

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

    #[test]
    fn test_hmac_deterministic() {
        let key = b"test-key-32-bytes-long-padding!!";
        let message = b"hello world";

        let hmac1 = compute_mac(key, message).unwrap();
        let hmac2 = compute_mac(key, message).unwrap();

        assert_eq!(hmac1, hmac2, "HMAC should be deterministic");
    }

    #[test]
    fn test_hmac_different_keys() {
        let key1 = b"key1-32-bytes-long-padding!!!!!!";
        let key2 = b"key2-32-bytes-long-padding!!!!!!";
        let message = b"hello world";

        let hmac1 = compute_mac(key1, message).unwrap();
        let hmac2 = compute_mac(key2, message).unwrap();

        assert_ne!(
            hmac1, hmac2,
            "Different keys should produce different HMACs"
        );
    }

    #[test]
    fn test_kdf_deterministic() {
        let mantra = "my secret server mantra";
        let params = KdfParams::new_v2([42u8; 32]);

        let key1 = KeyGenerator::derive_from_passphrase(mantra, &params).unwrap();
        let key2 = KeyGenerator::derive_from_passphrase(mantra, &params).unwrap();

        assert_eq!(
            key1, key2,
            "Derivation should be exactly deterministic for the same mantra and params"
        );
    }

    #[test]
    fn test_kdf_different_salts() {
        let mantra = "my secret server mantra";
        let params1 = KdfParams::new_v2([1u8; 32]);
        let params2 = KdfParams::new_v2([2u8; 32]);

        let key1 = KeyGenerator::derive_from_passphrase(mantra, &params1).unwrap();
        let key2 = KeyGenerator::derive_from_passphrase(mantra, &params2).unwrap();

        assert_ne!(key1, key2, "Different salts must produce different keys");
    }

    #[test]
    fn test_kdf_different_mantras() {
        let params = KdfParams::new_v2([42u8; 32]);

        let key1 = KeyGenerator::derive_from_passphrase("mantra A", &params).unwrap();
        let key2 = KeyGenerator::derive_from_passphrase("mantra B", &params).unwrap();

        assert_ne!(key1, key2, "Different mantras must produce different keys");
    }

    #[test]
    fn test_kdf_nfkc_normalization() {
        // "é" can be represented as a single code point (U+00E9) or two (U+0065, U+0301)
        let composed = "stréssed";
        let decomposed = "stre\u{0301}ssed";

        let params = KdfParams::new_v2([42u8; 32]);

        let key1 = KeyGenerator::derive_from_passphrase(composed, &params).unwrap();
        let key2 = KeyGenerator::derive_from_passphrase(decomposed, &params).unwrap();

        assert_eq!(
            key1, key2,
            "NFKC normalization must ensure visually identical mantras yield the same key"
        );

        // Also test leading/trailing whitespace trimming
        let trimmed = KeyGenerator::derive_from_passphrase("  stréssed \n ", &params).unwrap();
        assert_eq!(key1, trimmed, "Whitespace should be trimmed securely");
    }
}