//! Key derivation functions for Synor.
//!
//! Provides secure key derivation using:
//! - HKDF (HMAC-based Key Derivation Function) for deriving keys from high-entropy seeds
//! - PBKDF2 for password-based key derivation
//! - BIP-32 style hierarchical deterministic key derivation

use hkdf::Hkdf;
use hmac::Hmac;
use pbkdf2::pbkdf2_hmac;
use sha3::Sha3_256;
use thiserror::Error;
use zeroize::ZeroizeOnDrop;

/// Domain separation tags for different key types
pub mod domain {
    pub const ED25519_KEY: &[u8] = b"synor/ed25519/v1";
    pub const DILITHIUM_KEY: &[u8] = b"synor/dilithium/v1";
    pub const ENCRYPTION_KEY: &[u8] = b"synor/encryption/v1";
    pub const CHILD_KEY: &[u8] = b"synor/child/v1";
}

/// Derives a key from a seed using HKDF-SHA3-256.
///
/// # Arguments
/// * `seed` - The high-entropy input key material (IKM)
/// * `salt` - Optional salt value (can be empty for most uses)
/// * `info` - Context/application-specific info for domain separation
/// * `output_len` - Desired output key length
///
/// # Returns
/// A vector of derived key bytes
pub fn derive_key(
    seed: &[u8],
    salt: &[u8],
    info: &[u8],
    output_len: usize,
) -> Result<DerivedKey, DeriveKeyError> {
    if seed.is_empty() {
        return Err(DeriveKeyError::EmptySeed);
    }
    if output_len == 0 || output_len > 255 * 32 {
        return Err(DeriveKeyError::InvalidOutputLength(output_len));
    }

    let hk = Hkdf::<Sha3_256>::new(Some(salt), seed);
    let mut output = vec![0u8; output_len];
    hk.expand(info, &mut output)
        .map_err(|_| DeriveKeyError::ExpansionFailed)?;

    Ok(DerivedKey(output))
}

/// Derives a 32-byte key for Ed25519 signing.
pub fn derive_ed25519_key(master_seed: &[u8], account: u32) -> Result<DerivedKey, DeriveKeyError> {
    let info = [domain::ED25519_KEY, &account.to_be_bytes()].concat();
    derive_key(master_seed, b"", &info, 32)
}

/// Derives a 32-byte seed for Dilithium key generation.
pub fn derive_dilithium_seed(
    master_seed: &[u8],
    account: u32,
) -> Result<DerivedKey, DeriveKeyError> {
    let info = [domain::DILITHIUM_KEY, &account.to_be_bytes()].concat();
    derive_key(master_seed, b"", &info, 32)
}

/// Derives a 32-byte encryption key.
pub fn derive_encryption_key(
    master_seed: &[u8],
    purpose: &[u8],
) -> Result<DerivedKey, DeriveKeyError> {
    let info = [domain::ENCRYPTION_KEY, purpose].concat();
    derive_key(master_seed, b"", &info, 32)
}

/// Derives a child key from a parent key using BIP-32 style derivation.
pub fn derive_child_key(
    parent_key: &[u8],
    chain_code: &[u8; 32],
    index: u32,
) -> Result<(DerivedKey, [u8; 32]), DeriveKeyError> {
    use hmac::Mac;

    if parent_key.len() != 32 {
        return Err(DeriveKeyError::InvalidKeyLength);
    }

    let mut hmac =
        Hmac::<Sha3_256>::new_from_slice(chain_code).map_err(|_| DeriveKeyError::HmacError)?;

    // For hardened keys (index >= 0x80000000), use parent key
    // For normal keys, use parent public key (not implemented here for simplicity)
    if index >= 0x80000000 {
        hmac.update(&[0x00]);
        hmac.update(parent_key);
    } else {
        // Normal derivation - would need public key here
        // For now, treat all as hardened
        hmac.update(&[0x00]);
        hmac.update(parent_key);
    }
    hmac.update(&index.to_be_bytes());

    let result = hmac.finalize().into_bytes();

    // Split into child key and new chain code
    let mut child_key = [0u8; 32];
    let mut new_chain_code = [0u8; 32];

    // Use HKDF to expand to 64 bytes
    let hk = Hkdf::<Sha3_256>::new(None, &result);
    let mut expanded = [0u8; 64];
    hk.expand(domain::CHILD_KEY, &mut expanded)
        .map_err(|_| DeriveKeyError::ExpansionFailed)?;

    child_key.copy_from_slice(&expanded[..32]);
    new_chain_code.copy_from_slice(&expanded[32..]);

    Ok((DerivedKey(child_key.to_vec()), new_chain_code))
}

/// Derives a key from a password using PBKDF2-HMAC-SHA3-256.
///
/// # Arguments
/// * `password` - The password to derive from
/// * `salt` - A unique salt for this derivation
/// * `iterations` - Number of PBKDF2 iterations (recommended: 100,000+)
/// * `output_len` - Desired output key length
pub fn derive_from_password(
    password: &[u8],
    salt: &[u8],
    iterations: u32,
    output_len: usize,
) -> Result<DerivedKey, DeriveKeyError> {
    if password.is_empty() {
        return Err(DeriveKeyError::EmptyPassword);
    }
    if salt.len() < 8 {
        return Err(DeriveKeyError::SaltTooShort);
    }
    if iterations < 10_000 {
        return Err(DeriveKeyError::InsufficientIterations);
    }
    if output_len == 0 || output_len > 64 {
        return Err(DeriveKeyError::InvalidOutputLength(output_len));
    }

    let mut output = vec![0u8; output_len];
    pbkdf2_hmac::<Sha3_256>(password, salt, iterations, &mut output);

    Ok(DerivedKey(output))
}

/// BIP-44 style derivation path for Synor.
/// Path: m/44'/synor'/account'/change/index
/// Synor coin type: 0x5359 (SY in ASCII)
pub struct DerivationPath {
    /// Account number (hardened)
    pub account: u32,
    /// Change: 0 = external, 1 = internal
    pub change: u32,
    /// Address index
    pub index: u32,
}

impl DerivationPath {
    /// Synor coin type for BIP-44 (0x5359 = "SY")
    pub const COIN_TYPE: u32 = 0x5359;

    /// Creates a new derivation path.
    pub fn new(account: u32, change: u32, index: u32) -> Self {
        DerivationPath {
            account,
            change,
            index,
        }
    }

    /// Creates the default path (first external address).
    pub fn default_external() -> Self {
        DerivationPath {
            account: 0,
            change: 0,
            index: 0,
        }
    }

    /// Creates a path for the nth external address.
    pub fn external(account: u32, index: u32) -> Self {
        DerivationPath {
            account,
            change: 0,
            index,
        }
    }

    /// Creates a path for the nth internal (change) address.
    pub fn internal(account: u32, index: u32) -> Self {
        DerivationPath {
            account,
            change: 1,
            index,
        }
    }
}

impl Default for DerivationPath {
    fn default() -> Self {
        Self::default_external()
    }
}

impl std::fmt::Display for DerivationPath {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "m/44'/{}'/{}'/{}/{}",
            Self::COIN_TYPE,
            self.account,
            self.change,
            self.index
        )
    }
}

/// A derived key that automatically zeroizes on drop.
#[derive(Clone, ZeroizeOnDrop)]
pub struct DerivedKey(Vec<u8>);

impl DerivedKey {
    /// Returns the key bytes.
    pub fn as_bytes(&self) -> &[u8] {
        &self.0
    }

    /// Consumes the key and returns the bytes.
    pub fn into_bytes(mut self) -> Vec<u8> {
        std::mem::take(&mut self.0)
    }

    /// Returns the key length.
    pub fn len(&self) -> usize {
        self.0.len()
    }

    /// Returns true if the key is empty.
    pub fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    /// Converts to a fixed-size array.
    pub fn to_array<const N: usize>(&self) -> Option<[u8; N]> {
        if self.0.len() != N {
            return None;
        }
        let mut arr = [0u8; N];
        arr.copy_from_slice(&self.0);
        Some(arr)
    }
}

impl AsRef<[u8]> for DerivedKey {
    fn as_ref(&self) -> &[u8] {
        &self.0
    }
}

impl std::fmt::Debug for DerivedKey {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "DerivedKey([{} bytes])", self.0.len())
    }
}

/// Errors that can occur during key derivation.
#[derive(Debug, Error)]
pub enum DeriveKeyError {
    #[error("Seed cannot be empty")]
    EmptySeed,

    #[error("Password cannot be empty")]
    EmptyPassword,

    #[error("Salt must be at least 8 bytes")]
    SaltTooShort,

    #[error("At least 10,000 iterations required")]
    InsufficientIterations,

    #[error("Invalid output length: {0}")]
    InvalidOutputLength(usize),

    #[error("Invalid key length")]
    InvalidKeyLength,

    #[error("HKDF expansion failed")]
    ExpansionFailed,

    #[error("HMAC operation failed")]
    HmacError,
}

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

    #[test]
    fn test_derive_key() {
        let seed = b"test seed for key derivation";
        let key = derive_key(seed, b"salt", b"info", 32).unwrap();
        assert_eq!(key.len(), 32);

        // Same inputs should produce same output
        let key2 = derive_key(seed, b"salt", b"info", 32).unwrap();
        assert_eq!(key.as_bytes(), key2.as_bytes());

        // Different info should produce different output
        let key3 = derive_key(seed, b"salt", b"different", 32).unwrap();
        assert_ne!(key.as_bytes(), key3.as_bytes());
    }

    #[test]
    fn test_derive_ed25519_key() {
        let master_seed = [42u8; 64];
        let key0 = derive_ed25519_key(&master_seed, 0).unwrap();
        let key1 = derive_ed25519_key(&master_seed, 1).unwrap();

        assert_eq!(key0.len(), 32);
        assert_eq!(key1.len(), 32);
        assert_ne!(key0.as_bytes(), key1.as_bytes());
    }

    #[test]
    fn test_derive_from_password() {
        let password = b"correct horse battery staple";
        let salt = b"unique salt value";

        let key = derive_from_password(password, salt, 100_000, 32).unwrap();
        assert_eq!(key.len(), 32);

        // Same password/salt should produce same key
        let key2 = derive_from_password(password, salt, 100_000, 32).unwrap();
        assert_eq!(key.as_bytes(), key2.as_bytes());

        // Different password should produce different key
        let key3 = derive_from_password(b"different password", salt, 100_000, 32).unwrap();
        assert_ne!(key.as_bytes(), key3.as_bytes());
    }

    #[test]
    fn test_derive_child_key() {
        let parent = [1u8; 32];
        let chain_code = [2u8; 32];

        let (child0, cc0) = derive_child_key(&parent, &chain_code, 0x80000000).unwrap();
        let (child1, cc1) = derive_child_key(&parent, &chain_code, 0x80000001).unwrap();

        assert_eq!(child0.len(), 32);
        assert_ne!(child0.as_bytes(), child1.as_bytes());
        assert_ne!(cc0, cc1);
    }

    #[test]
    fn test_derivation_path() {
        let path = DerivationPath::default_external();
        assert_eq!(path.to_string(), "m/44'/21337'/0'/0/0");

        let path2 = DerivationPath::external(1, 5);
        assert_eq!(path2.to_string(), "m/44'/21337'/1'/0/5");

        let path3 = DerivationPath::internal(0, 3);
        assert_eq!(path3.to_string(), "m/44'/21337'/0'/1/3");
    }

    #[test]
    fn test_empty_seed_error() {
        let result = derive_key(&[], b"salt", b"info", 32);
        assert!(matches!(result, Err(DeriveKeyError::EmptySeed)));
    }

    #[test]
    fn test_password_validation() {
        // Empty password
        assert!(matches!(
            derive_from_password(&[], b"long enough salt", 100_000, 32),
            Err(DeriveKeyError::EmptyPassword)
        ));

        // Salt too short
        assert!(matches!(
            derive_from_password(b"password", b"short", 100_000, 32),
            Err(DeriveKeyError::SaltTooShort)
        ));

        // Not enough iterations
        assert!(matches!(
            derive_from_password(b"password", b"long enough salt", 1000, 32),
            Err(DeriveKeyError::InsufficientIterations)
        ));
    }
}