//! Key Manager
//!
//! Manages the lifecycle of cryptographic keys:
//! - Key generation with secure random
//! - Session key derivation
//! - Sync token creation and verification
//! - Key rotation and revocation
//! - Automatic memory zeroing of key material

use std::collections::HashMap;
use std::time::{SystemTime, UNIX_EPOCH};

use rand::RngCore;
use serde::{Deserialize, Serialize};
use zeroize::{Zeroize, ZeroizeOnDrop};

use crate::error::{CryptoError, CryptoResult};
use crate::keys::hkdf::{derive_key_32, derive_sync_token};
use crate::keys::hmac::{create_sync_token as create_hmac_token, verify_sync_token as verify_hmac_token};

/// A cryptographic key with metadata
#[derive(Clone, ZeroizeOnDrop)]
pub struct PhantomKey {
    /// Unique identifier for this key
    #[zeroize(skip)]
    pub key_id: String,
    /// The actual key material (zeroed on drop)
    key_material: [u8; 32],
    /// Unix timestamp when key was created
    #[zeroize(skip)]
    pub created_at: u64,
    /// Unix timestamp when key expires (None = never)
    #[zeroize(skip)]
    pub expires_at: Option<u64>,
    /// Whether this key has been revoked
    #[zeroize(skip)]
    pub revoked: bool,
    /// Optional channel assignment
    #[zeroize(skip)]
    pub channel: Option<String>,
}

impl PhantomKey {
    /// Check if the key is currently valid
    pub fn is_valid(&self) -> bool {
        if self.revoked {
            return false;
        }

        if let Some(expires) = self.expires_at {
            let now = current_timestamp();
            if now >= expires {
                return false;
            }
        }

        true
    }

    /// Get the key material (only if valid)
    pub fn material(&self) -> CryptoResult<&[u8; 32]> {
        if self.revoked {
            return Err(CryptoError::KeyRevoked(self.key_id.clone()));
        }

        if let Some(expires) = self.expires_at {
            let now = current_timestamp();
            if now >= expires {
                return Err(CryptoError::KeyExpired(self.key_id.clone()));
            }
        }

        Ok(&self.key_material)
    }

    /// Get time until expiration in seconds (None if no expiration)
    pub fn time_to_expiry(&self) -> Option<u64> {
        self.expires_at.map(|exp| {
            let now = current_timestamp();
            exp.saturating_sub(now)
        })
    }
}

/// Key metadata for serialization (without key material)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct KeyMetadata {
    /// Key identifier
    pub key_id: String,
    /// Creation timestamp
    pub created_at: u64,
    /// Expiration timestamp
    pub expires_at: Option<u64>,
    /// Revocation status
    pub revoked: bool,
    /// Channel assignment
    pub channel: Option<String>,
}

impl From<&PhantomKey> for KeyMetadata {
    fn from(key: &PhantomKey) -> Self {
        Self {
            key_id: key.key_id.clone(),
            created_at: key.created_at,
            expires_at: key.expires_at,
            revoked: key.revoked,
            channel: key.channel.clone(),
        }
    }
}

/// Key Manager
///
/// Manages all cryptographic keys for covert channel operations.
pub struct KeyManager {
    keys: HashMap<String, PhantomKey>,
    /// Prefix for generated key IDs
    key_prefix: String,
    /// Default key lifetime in seconds (None = no expiration)
    default_lifetime: Option<u64>,
}

impl Default for KeyManager {
    fn default() -> Self {
        Self::new()
    }
}

impl KeyManager {
    /// Create a new key manager
    pub fn new() -> Self {
        Self {
            keys: HashMap::new(),
            key_prefix: "phantom".to_string(),
            default_lifetime: None,
        }
    }

    /// Create a key manager with custom settings
    pub fn with_config(prefix: impl Into<String>, default_lifetime_secs: Option<u64>) -> Self {
        Self {
            keys: HashMap::new(),
            key_prefix: prefix.into(),
            default_lifetime: default_lifetime_secs,
        }
    }

    /// Generate a new key
    ///
    /// # Returns
    /// The key ID of the newly generated key
    pub fn generate_key(&mut self) -> CryptoResult<String> {
        self.generate_key_with_lifetime(self.default_lifetime)
    }

    /// Generate a new key with a specific lifetime
    ///
    /// # Arguments
    /// * `lifetime_secs` - Key lifetime in seconds (None = no expiration)
    pub fn generate_key_with_lifetime(&mut self, lifetime_secs: Option<u64>) -> CryptoResult<String> {
        let mut key_material = [0u8; 32];
        let mut rng = rand::thread_rng();
        rng.fill_bytes(&mut key_material);

        let now = current_timestamp();
        let key_id = generate_key_id(&self.key_prefix);

        let key = PhantomKey {
            key_id: key_id.clone(),
            key_material,
            created_at: now,
            expires_at: lifetime_secs.map(|l| now + l),
            revoked: false,
            channel: None,
        };

        self.keys.insert(key_id.clone(), key);
        Ok(key_id)
    }

    /// Generate a key for a specific channel
    pub fn generate_channel_key(&mut self, channel: &str) -> CryptoResult<String> {
        let key_id = self.generate_key()?;
        if let Some(key) = self.keys.get_mut(&key_id) {
            key.channel = Some(channel.to_string());
        }
        Ok(key_id)
    }

    /// Import an existing key
    ///
    /// # Arguments
    /// * `key_id` - Identifier for the key
    /// * `key_material` - The 32-byte key
    /// * `expires_at` - Optional expiration timestamp
    pub fn import_key(
        &mut self,
        key_id: impl Into<String>,
        key_material: [u8; 32],
        expires_at: Option<u64>,
    ) -> CryptoResult<()> {
        let key_id = key_id.into();

        if self.keys.contains_key(&key_id) {
            return Err(CryptoError::KeyDerivationFailed(
                format!("Key {} already exists", key_id)
            ));
        }

        let key = PhantomKey {
            key_id: key_id.clone(),
            key_material,
            created_at: current_timestamp(),
            expires_at,
            revoked: false,
            channel: None,
        };

        self.keys.insert(key_id, key);
        Ok(())
    }

    /// Get a key by ID
    pub fn get_key(&self, key_id: &str) -> CryptoResult<&PhantomKey> {
        self.keys.get(key_id)
            .ok_or_else(|| CryptoError::KeyNotFound(key_id.to_string()))
    }

    /// Get key material by ID (only if valid)
    pub fn get_key_material(&self, key_id: &str) -> CryptoResult<&[u8; 32]> {
        let key = self.get_key(key_id)?;
        key.material()
    }

    /// Check if a key exists and is valid
    pub fn is_key_valid(&self, key_id: &str) -> bool {
        self.keys.get(key_id)
            .map(|k| k.is_valid())
            .unwrap_or(false)
    }

    /// Derive a session key from a master key
    ///
    /// # Arguments
    /// * `key_id` - The master key ID
    /// * `session_id` - Unique session identifier
    ///
    /// # Returns
    /// 32-byte derived session key
    pub fn derive_session_key(&self, key_id: &str, session_id: &str) -> CryptoResult<[u8; 32]> {
        let master = self.get_key_material(key_id)?;
        let channel = self.get_key(key_id)?
            .channel.as_deref()
            .unwrap_or("default");

        let info = format!("phantom:session:{}:{}", channel, session_id);
        derive_key_32(master, None, info.as_bytes())
    }

    /// Create a sync token for agent synchronization
    ///
    /// # Arguments
    /// * `key_id` - The shared key ID
    /// * `timestamp` - Unix timestamp for token validity
    ///
    /// # Returns
    /// 16-byte sync token
    pub fn create_sync_token(&self, key_id: &str, timestamp: u64) -> CryptoResult<[u8; 16]> {
        let key = self.get_key_material(key_id)?;
        derive_sync_token(key, timestamp, &["default"])
    }

    /// Verify a sync token
    ///
    /// # Arguments
    /// * `key_id` - The shared key ID
    /// * `token` - The token to verify
    /// * `timestamp` - The claimed timestamp
    /// * `tolerance_secs` - Allowed time drift in seconds
    pub fn verify_sync_token(
        &self,
        key_id: &str,
        token: &[u8],
        timestamp: u64,
        tolerance_secs: u64,
    ) -> CryptoResult<bool> {
        let key = self.get_key_material(key_id)?;

        // Check within tolerance window
        for offset in 0..=tolerance_secs {
            let expected = derive_sync_token(key, timestamp + offset, &["default"])?;
            if expected == token {
                return Ok(true);
            }

            if offset > 0 && timestamp >= offset {
                let expected = derive_sync_token(key, timestamp - offset, &["default"])?;
                if expected == token {
                    return Ok(true);
                }
            }
        }

        Ok(false)
    }

    /// Rotate a key (generate new key, mark old as expired)
    ///
    /// # Arguments
    /// * `key_id` - The key to rotate
    /// * `grace_period_secs` - How long the old key remains valid
    ///
    /// # Returns
    /// The new key ID
    pub fn rotate_key(&mut self, key_id: &str, grace_period_secs: u64) -> CryptoResult<String> {
        let old_key = self.keys.get_mut(key_id)
            .ok_or_else(|| CryptoError::KeyNotFound(key_id.to_string()))?;

        if old_key.revoked {
            return Err(CryptoError::KeyRevoked(key_id.to_string()));
        }

        // Set expiration on old key
        let now = current_timestamp();
        old_key.expires_at = Some(now + grace_period_secs);

        // Generate new key with same channel
        let channel = old_key.channel.clone();
        let new_key_id = self.generate_key()?;

        if let (Some(channel), Some(new_key)) = (channel, self.keys.get_mut(&new_key_id)) {
            new_key.channel = Some(channel);
        }

        Ok(new_key_id)
    }

    /// Revoke a key immediately
    ///
    /// Revoked keys cannot be used for any operations.
    pub fn revoke_key(&mut self, key_id: &str) -> CryptoResult<()> {
        let key = self.keys.get_mut(key_id)
            .ok_or_else(|| CryptoError::KeyNotFound(key_id.to_string()))?;

        key.revoked = true;
        Ok(())
    }

    /// Remove a key entirely (use with caution)
    pub fn delete_key(&mut self, key_id: &str) -> CryptoResult<()> {
        self.keys.remove(key_id)
            .map(|_| ())
            .ok_or_else(|| CryptoError::KeyNotFound(key_id.to_string()))
    }

    /// Get metadata for all keys
    pub fn list_keys(&self) -> Vec<KeyMetadata> {
        self.keys.values()
            .map(KeyMetadata::from)
            .collect()
    }

    /// Get metadata for valid keys only
    pub fn list_valid_keys(&self) -> Vec<KeyMetadata> {
        self.keys.values()
            .filter(|k| k.is_valid())
            .map(KeyMetadata::from)
            .collect()
    }

    /// Get number of keys
    pub fn key_count(&self) -> usize {
        self.keys.len()
    }

    /// Clean up expired and revoked keys
    pub fn cleanup(&mut self) -> usize {
        let before = self.keys.len();
        self.keys.retain(|_, k| k.is_valid());
        before - self.keys.len()
    }

    /// Get keys that will expire within the given time
    pub fn expiring_keys(&self, within_secs: u64) -> Vec<KeyMetadata> {
        let now = current_timestamp();
        let threshold = now + within_secs;

        self.keys.values()
            .filter(|k| {
                k.is_valid() &&
                k.expires_at.map(|e| e <= threshold).unwrap_or(false)
            })
            .map(KeyMetadata::from)
            .collect()
    }
}

// Helper functions

fn current_timestamp() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap()
        .as_secs()
}

fn generate_key_id(prefix: &str) -> String {
    let mut random = [0u8; 8];
    rand::thread_rng().fill_bytes(&mut random);
    format!("{}-{}", prefix, hex::encode(random))
}

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

    #[test]
    fn test_generate_key() {
        let mut manager = KeyManager::new();
        let key_id = manager.generate_key().unwrap();

        assert!(key_id.starts_with("phantom-"));
        assert!(manager.is_key_valid(&key_id));
    }

    #[test]
    fn test_key_material() {
        let mut manager = KeyManager::new();
        let key_id = manager.generate_key().unwrap();

        let material = manager.get_key_material(&key_id).unwrap();
        assert_eq!(material.len(), 32);
    }

    #[test]
    fn test_key_not_found() {
        let manager = KeyManager::new();
        let result = manager.get_key("nonexistent");
        assert!(matches!(result, Err(CryptoError::KeyNotFound(_))));
    }

    #[test]
    fn test_key_with_expiration() {
        let mut manager = KeyManager::new();
        let key_id = manager.generate_key_with_lifetime(Some(3600)).unwrap();

        let key = manager.get_key(&key_id).unwrap();
        assert!(key.expires_at.is_some());
        assert!(key.is_valid());
    }

    #[test]
    fn test_channel_key() {
        let mut manager = KeyManager::new();
        let key_id = manager.generate_channel_key("TCOE").unwrap();

        let key = manager.get_key(&key_id).unwrap();
        assert_eq!(key.channel.as_deref(), Some("TCOE"));
    }

    #[test]
    fn test_import_key() {
        let mut manager = KeyManager::new();
        let material = [0xABu8; 32];

        manager.import_key("imported-key", material, None).unwrap();

        let retrieved = manager.get_key_material("imported-key").unwrap();
        assert_eq!(*retrieved, material);
    }

    #[test]
    fn test_import_duplicate() {
        let mut manager = KeyManager::new();
        let material = [0xABu8; 32];

        manager.import_key("dup-key", material, None).unwrap();
        let result = manager.import_key("dup-key", material, None);
        assert!(result.is_err());
    }

    #[test]
    fn test_derive_session_key() {
        let mut manager = KeyManager::new();
        let key_id = manager.generate_key().unwrap();

        let session1 = manager.derive_session_key(&key_id, "session-1").unwrap();
        let session2 = manager.derive_session_key(&key_id, "session-2").unwrap();

        assert_eq!(session1.len(), 32);
        assert_ne!(session1, session2);

        // Same session ID produces same key
        let session1_again = manager.derive_session_key(&key_id, "session-1").unwrap();
        assert_eq!(session1, session1_again);
    }

    #[test]
    fn test_sync_token() {
        let mut manager = KeyManager::new();
        let key_id = manager.generate_key().unwrap();

        let timestamp = current_timestamp();
        let token = manager.create_sync_token(&key_id, timestamp).unwrap();

        assert_eq!(token.len(), 16);

        // Verify exact timestamp
        assert!(manager.verify_sync_token(&key_id, &token, timestamp, 0).unwrap());

        // Verify with tolerance
        assert!(manager.verify_sync_token(&key_id, &token, timestamp + 1, 5).unwrap());
        assert!(manager.verify_sync_token(&key_id, &token, timestamp - 1, 5).unwrap());

        // Wrong timestamp outside tolerance
        assert!(!manager.verify_sync_token(&key_id, &token, timestamp + 100, 5).unwrap());
    }

    #[test]
    fn test_revoke_key() {
        let mut manager = KeyManager::new();
        let key_id = manager.generate_key().unwrap();

        assert!(manager.is_key_valid(&key_id));

        manager.revoke_key(&key_id).unwrap();

        assert!(!manager.is_key_valid(&key_id));

        // Trying to use revoked key should fail
        let result = manager.get_key_material(&key_id);
        assert!(matches!(result, Err(CryptoError::KeyRevoked(_))));
    }

    #[test]
    fn test_rotate_key() {
        let mut manager = KeyManager::new();
        let old_key_id = manager.generate_channel_key("TCOE").unwrap();

        let new_key_id = manager.rotate_key(&old_key_id, 3600).unwrap();

        // Old key should still be valid (grace period)
        assert!(manager.is_key_valid(&old_key_id));

        // New key should be valid
        assert!(manager.is_key_valid(&new_key_id));

        // New key should have same channel
        let new_key = manager.get_key(&new_key_id).unwrap();
        assert_eq!(new_key.channel.as_deref(), Some("TCOE"));
    }

    #[test]
    fn test_delete_key() {
        let mut manager = KeyManager::new();
        let key_id = manager.generate_key().unwrap();

        assert_eq!(manager.key_count(), 1);

        manager.delete_key(&key_id).unwrap();

        assert_eq!(manager.key_count(), 0);
        assert!(!manager.is_key_valid(&key_id));
    }

    #[test]
    fn test_list_keys() {
        let mut manager = KeyManager::new();
        manager.generate_key().unwrap();
        manager.generate_key().unwrap();
        manager.generate_key().unwrap();

        let keys = manager.list_keys();
        assert_eq!(keys.len(), 3);
    }

    #[test]
    fn test_list_valid_keys() {
        let mut manager = KeyManager::new();
        let key1 = manager.generate_key().unwrap();
        let key2 = manager.generate_key().unwrap();
        manager.generate_key().unwrap();

        manager.revoke_key(&key1).unwrap();
        manager.revoke_key(&key2).unwrap();

        let valid_keys = manager.list_valid_keys();
        assert_eq!(valid_keys.len(), 1);
    }

    #[test]
    fn test_key_metadata() {
        let mut manager = KeyManager::new();
        let key_id = manager.generate_channel_key("EGE").unwrap();

        let key = manager.get_key(&key_id).unwrap();
        let metadata: KeyMetadata = key.into();

        assert_eq!(metadata.key_id, key_id);
        assert_eq!(metadata.channel.as_deref(), Some("EGE"));
        assert!(!metadata.revoked);
    }

    #[test]
    fn test_custom_config() {
        let mut manager = KeyManager::with_config("pqc", Some(7200));
        let key_id = manager.generate_key().unwrap();

        assert!(key_id.starts_with("pqc-"));

        let key = manager.get_key(&key_id).unwrap();
        assert!(key.expires_at.is_some());
    }

    #[test]
    fn test_zeroize_on_drop() {
        let mut manager = KeyManager::new();
        let key_id = manager.generate_key().unwrap();

        // Get a copy of the key material
        let material_copy = *manager.get_key_material(&key_id).unwrap();

        // Delete the key (should zero the material)
        manager.delete_key(&key_id).unwrap();

        // We can't directly verify zeroing, but we can verify the key is gone
        assert!(manager.get_key(&key_id).is_err());

        // The copy we made should still have the original bytes
        assert_ne!(material_copy, [0u8; 32]);
    }
}
