//! Checkpoint cryptographic signature system
//!
//! Provides HMAC-SHA256 signing and verification for ML model checkpoints
//! to prevent tampering and verify authenticity.

use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;

use chrono::{DateTime, Datelike, Utc};
use hmac::{Hmac, Mac};
use sha2::Sha256;
use tokio::sync::RwLock;
use tracing::{debug, error, info, warn};

use crate::ModelType;
use crate::MLError;

type HmacSha256 = Hmac<Sha256>;

/// Signature information for a checkpoint
#[derive(Debug, Clone)]
pub struct SignatureInfo {
    /// HMAC-SHA256 signature (hex-encoded)
    pub signature: String,
    /// Signing algorithm identifier
    pub algorithm: String,
    /// Key ID for rotation support
    pub key_id: String,
    /// Signature timestamp
    pub signed_at: DateTime<Utc>,
}

/// Checkpoint signer with Vault integration
///
/// Provides cryptographic signatures for checkpoints using HMAC-SHA256.
/// Keys are stored in Vault with support for rotation and per-model-type keys.
#[derive(Clone)]
pub struct CheckpointSigner {
    /// In-memory key cache to avoid excessive Vault calls
    key_cache: Arc<RwLock<KeyCache>>,
    /// Cache TTL (default: 5 minutes)
    cache_ttl: Duration,
}

impl std::fmt::Debug for CheckpointSigner {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("CheckpointSigner")
            .field("cache_ttl", &self.cache_ttl)
            .field("key_cache", &"Arc<RwLock<KeyCache>>")
            .finish()
    }
}

/// Key cache with TTL support
#[derive(Debug)]
struct KeyCache {
    keys: HashMap<String, CachedKey>,
}

#[derive(Debug, Clone)]
struct CachedKey {
    key_data: Vec<u8>,
    cached_at: DateTime<Utc>,
}

impl CheckpointSigner {
    /// Create a new checkpoint signer
    ///
    /// # Arguments
    /// * `cache_ttl` - Cache TTL for Vault keys (default: 5 minutes)
    pub fn new(cache_ttl: Option<Duration>) -> Self {
        Self {
            key_cache: Arc::new(RwLock::new(KeyCache {
                keys: HashMap::new(),
            })),
            cache_ttl: cache_ttl.unwrap_or(Duration::from_secs(300)),
        }
    }

    /// Sign checkpoint data with HMAC-SHA256
    pub async fn sign_checkpoint(
        &self,
        data: &[u8],
        model_type: ModelType,
    ) -> Result<SignatureInfo, MLError> {
        let start = std::time::Instant::now();

        // Generate current key ID (format: "2024-Q4")
        let key_id = self.generate_key_id();

        // Get signing key (from cache or Vault)
        let signing_key = self.get_signing_key(&key_id, model_type).await?;

        // Compute HMAC-SHA256 signature
        let mut mac =
            HmacSha256::new_from_slice(&signing_key).map_err(|e| MLError::ValidationError {
                message: format!("Failed to create HMAC instance: {}", e),
            })?;

        mac.update(data);
        let signature_bytes = mac.finalize().into_bytes();
        let signature = hex::encode(signature_bytes);

        let duration = start.elapsed();
        debug!(
            "Signed checkpoint for {:?} in {:?} (signature: {}...)",
            model_type,
            duration,
            signature.get(..8).unwrap_or(&signature)
        );

        Ok(SignatureInfo {
            signature,
            algorithm: "HMAC-SHA256".to_owned(),
            key_id,
            signed_at: Utc::now(),
        })
    }

    /// Verify checkpoint signature
    ///
    /// Uses constant-time comparison to prevent timing attacks
    pub async fn verify_signature(
        &self,
        data: &[u8],
        signature: &str,
        key_id: &str,
        model_type: ModelType,
    ) -> Result<(), MLError> {
        let start = std::time::Instant::now();

        // Get signing key (from cache or Vault)
        let signing_key = self.get_signing_key_by_id(key_id, model_type).await?;

        // Compute expected signature
        let mut mac =
            HmacSha256::new_from_slice(&signing_key).map_err(|e| MLError::ValidationError {
                message: format!("Failed to create HMAC instance: {}", e),
            })?;

        mac.update(data);

        // Decode provided signature
        let signature_bytes = hex::decode(signature).map_err(|e| MLError::ValidationError {
            message: format!("Invalid signature hex encoding: {}", e),
        })?;

        // Constant-time comparison
        mac.verify_slice(&signature_bytes).map_err(|e| {
            error!(
                "Checkpoint signature verification failed - possible tampering (key_id: {}): {}",
                key_id, e
            );
            MLError::ValidationError {
                message: format!(
                    "Checkpoint signature verification failed - possible tampering (key_id: {}): {}",
                    key_id, e
                ),
            }
        })?;

        let duration = start.elapsed();
        debug!(
            "Verified checkpoint signature for {:?} in {:?} (key_id: {})",
            model_type, duration, key_id
        );

        Ok(())
    }

    /// Get signing key from cache or Vault
    async fn get_signing_key(
        &self,
        key_id: &str,
        model_type: ModelType,
    ) -> Result<Vec<u8>, MLError> {
        // Create composite cache key to differentiate between model types
        let cache_key = format!("{:?}-{}", model_type, key_id);

        // Check cache first
        {
            let cache = self.key_cache.read().await;
            if let Some(cached) = cache.keys.get(&cache_key) {
                let age = Utc::now()
                    .signed_duration_since(cached.cached_at)
                    .to_std()
                    .unwrap_or(Duration::MAX);

                if age < self.cache_ttl {
                    debug!("Key cache hit: {} (age: {:?})", cache_key, age);
                    return Ok(cached.key_data.clone());
                }
            }
        }

        // Cache miss or expired - fetch from Vault
        debug!("Key cache miss: {} - fetching from Vault", cache_key);
        let key_data = self.fetch_key_from_vault(key_id, model_type).await?;

        // Update cache
        {
            let mut cache = self.key_cache.write().await;
            cache.keys.insert(
                cache_key,
                CachedKey {
                    key_data: key_data.clone(),
                    cached_at: Utc::now(),
                },
            );
        }

        Ok(key_data)
    }

    /// Get signing key by specific ID (for verification of old checkpoints)
    async fn get_signing_key_by_id(
        &self,
        key_id: &str,
        model_type: ModelType,
    ) -> Result<Vec<u8>, MLError> {
        self.get_signing_key(key_id, model_type).await
    }

    /// Fetch signing key from Vault
    ///
    /// # Key Path Format
    /// `secret/foxhunt/ml/checkpoint-signing/{model_type}/{key_id}`
    async fn fetch_key_from_vault(
        &self,
        key_id: &str,
        model_type: ModelType,
    ) -> Result<Vec<u8>, MLError> {
        // For now, use environment variable as fallback
        // TODO: Replace with actual Vault integration when available
        let env_key = format!(
            "CHECKPOINT_SIGNING_KEY_{:?}_{}",
            model_type,
            key_id.replace('-', "_")
        );

        if let Ok(key_hex) = std::env::var(&env_key) {
            let key_data = hex::decode(&key_hex).map_err(|e| MLError::ConfigError(format!("Invalid key hex in {}: {}", env_key, e)))?;

            if key_data.len() != 32 {
                return Err(MLError::ConfigError(format!(
                        "Invalid key length in {} (expected 32 bytes, got {})",
                        env_key,
                        key_data.len()
                    )));
            }

            info!(
                "Loaded signing key for {:?} from environment (key_id: {})",
                model_type, key_id
            );
            Ok(key_data)
        } else {
            // Generate a default key for development/testing
            warn!(
                "No signing key found in Vault or environment for {:?} (key_id: {}), using default development key",
                model_type, key_id
            );

            // Generate deterministic key from model_type + key_id for consistency
            let seed = format!("foxhunt-{:?}-{}", model_type, key_id);
            let mut hasher = Sha256::new();
            use sha2::Digest;
            hasher.update(seed.as_bytes());
            let hash = hasher.finalize();
            Ok(hash.to_vec())
        }
    }

    /// Generate current key ID based on quarter
    ///
    /// # Format
    /// "YYYY-QN" (e.g., "2024-Q4")
    fn generate_key_id(&self) -> String {
        let now = Utc::now();
        let year = now.year();
        let quarter = (now.month() - 1) / 3 + 1;
        format!("{}-Q{}", year, quarter)
    }

    /// Rotate signing keys (quarterly maintenance task)
    pub async fn rotate_keys(&self) -> Result<(), MLError> {
        info!("Starting signing key rotation");

        // Generate next quarter's key ID
        let current_key_id = self.generate_key_id();
        let next_key_id = self.generate_next_key_id();

        info!("Current key: {}, Next key: {}", current_key_id, next_key_id);

        // In production, this would generate new keys in Vault
        warn!("Key rotation not fully implemented - would generate keys in Vault");

        // Clear cache to force reload
        {
            let mut cache = self.key_cache.write().await;
            cache.keys.clear();
        }

        info!("Signing key rotation completed");
        Ok(())
    }

    /// Generate next quarter's key ID
    fn generate_next_key_id(&self) -> String {
        let now = Utc::now();
        let year = now.year();
        let month = now.month();

        let (next_year, next_quarter) = if month >= 10 {
            (year + 1, 1)
        } else {
            let quarter = (month - 1) / 3 + 1;
            (year, quarter + 1)
        };

        format!("{}-Q{}", next_year, next_quarter)
    }

    /// Clear key cache (useful for testing)
    #[cfg(test)]
    pub async fn clear_cache(&self) {
        let mut cache = self.key_cache.write().await;
        cache.keys.clear();
    }
}

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

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

    #[tokio::test]
    async fn test_sign_and_verify_checkpoint() {
        let signer = CheckpointSigner::new(None);
        let data = b"test checkpoint data";

        // Sign checkpoint
        let sig_info = signer.sign_checkpoint(data, ModelType::DQN).await.unwrap();

        assert_eq!(sig_info.algorithm, "HMAC-SHA256");
        assert!(!sig_info.signature.is_empty());
        assert!(!sig_info.key_id.is_empty());

        // Verify signature
        let result = signer
            .verify_signature(data, &sig_info.signature, &sig_info.key_id, ModelType::DQN)
            .await;

        assert!(result.is_ok(), "Signature verification should succeed");
    }

    #[tokio::test]
    async fn test_verify_invalid_signature() {
        let signer = CheckpointSigner::new(None);
        let data = b"test checkpoint data";

        let sig_info = signer.sign_checkpoint(data, ModelType::DQN).await.unwrap();

        // Tamper with signature
        let mut tampered_sig = sig_info.signature.clone();
        tampered_sig.push_str("00");

        // Verification should fail
        let result = signer
            .verify_signature(data, &tampered_sig, &sig_info.key_id, ModelType::DQN)
            .await;

        assert!(result.is_err(), "Tampered signature should be rejected");
    }

    #[tokio::test]
    async fn test_verify_tampered_data() {
        let signer = CheckpointSigner::new(None);
        let data = b"test checkpoint data";

        let sig_info = signer.sign_checkpoint(data, ModelType::DQN).await.unwrap();

        // Tamper with data
        let tampered_data = b"test checkpoint DATA";

        // Verification should fail
        let result = signer
            .verify_signature(
                tampered_data,
                &sig_info.signature,
                &sig_info.key_id,
                ModelType::DQN,
            )
            .await;

        assert!(
            result.is_err(),
            "Signature should not verify with tampered data"
        );
    }

    #[tokio::test]
    async fn test_key_cache() {
        let signer = CheckpointSigner::new(Some(Duration::from_secs(60)));
        let data = b"test checkpoint data";

        // First sign (cache miss)
        let sig1 = signer.sign_checkpoint(data, ModelType::DQN).await.unwrap();

        // Second sign (cache hit)
        let sig2 = signer.sign_checkpoint(data, ModelType::DQN).await.unwrap();

        // Should use same key ID
        assert_eq!(sig1.key_id, sig2.key_id);

        // Signatures should be identical for same data
        assert_eq!(sig1.signature, sig2.signature);
    }

    #[tokio::test]
    async fn test_different_model_types() {
        let signer = CheckpointSigner::new(None);
        let data = b"test checkpoint data";

        let sig_dqn = signer.sign_checkpoint(data, ModelType::DQN).await.unwrap();

        let sig_ppo = signer.sign_checkpoint(data, ModelType::PPO).await.unwrap();

        // Different model types should produce different signatures
        // (due to different signing keys)
        assert_ne!(
            sig_dqn.signature, sig_ppo.signature,
            "Different model types should have different signatures"
        );
    }

    #[test]
    fn test_key_id_generation() {
        let signer = CheckpointSigner::new(None);
        let key_id = signer.generate_key_id();

        // Should match format "YYYY-QN"
        assert!(
            key_id.len() == 7 || key_id.len() == 8,
            "Key ID should be 7-8 characters"
        );
        assert!(key_id.contains("-Q"), "Key ID should contain quarter");
    }

    #[tokio::test]
    async fn test_signature_hex_encoding() {
        let signer = CheckpointSigner::new(None);
        let data = b"test checkpoint data";

        let sig_info = signer.sign_checkpoint(data, ModelType::DQN).await.unwrap();

        // Signature should be valid hex
        let decoded = hex::decode(&sig_info.signature);
        assert!(decoded.is_ok(), "Signature should be valid hex");
        assert_eq!(
            decoded.unwrap().len(),
            32,
            "HMAC-SHA256 should produce 32 bytes"
        );
    }
}
