//! SSH Key Management with KMS Integration
//!
//! Provides SSH key storage, retrieval, rotation, and audit logging
//! using the existing KMS infrastructure.

use super::{KmsBackend, KmsError, KeyData, KeyType, KeyMetadata, AuditEvent};
use async_trait::async_trait;
use chrono::{DateTime, Utc, Duration};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tokio::sync::RwLock;
use uuid::Uuid;

/// SSH key pair stored in KMS
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SshKeyPair {
    pub key_id: String,
    pub name: String,
    pub public_key: String,
    pub private_key_id: String,  // Reference to private key in KMS
    pub fingerprint: String,
    pub created_at: DateTime<Utc>,
    pub last_used: Option<DateTime<Utc>>,
    pub rotation_due: Option<DateTime<Utc>>,
    pub tags: Vec<String>,
    pub purpose: SshKeyPurpose,
}

/// Purpose of SSH key
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum SshKeyPurpose {
    ServerAccess,
    GitAccess,
    Automation,
    UserKey,
    ServiceAccount,
}

/// SSH key metadata for listing
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SshKeyMetadata {
    pub key_id: String,
    pub name: String,
    pub fingerprint: String,
    pub created_at: DateTime<Utc>,
    pub last_used: Option<DateTime<Utc>>,
    pub rotation_due: Option<DateTime<Utc>>,
    pub purpose: SshKeyPurpose,
}

/// SSH key access audit entry
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SshKeyAuditEntry {
    pub timestamp: DateTime<Utc>,
    pub key_id: String,
    pub action: SshKeyAction,
    pub user: Option<String>,
    pub ip_address: Option<String>,
    pub success: bool,
    pub error_message: Option<String>,
}

/// SSH key actions for audit
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum SshKeyAction {
    Created,
    Retrieved,
    Rotated,
    Deleted,
    Used,
    AuditViewed,
}

/// SSH key manager
pub struct SshKeyManager {
    kms_backend: Arc<dyn KmsBackend>,
    audit_log: Arc<RwLock<Vec<SshKeyAuditEntry>>>,
    rotation_policy: RotationPolicy,
}

/// Rotation policy for SSH keys
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RotationPolicy {
    pub enabled: bool,
    pub rotation_interval_days: u32,
    pub grace_period_days: u32,
    pub auto_rotate: bool,
}

impl Default for RotationPolicy {
    fn default() -> Self {
        Self {
            enabled: true,
            rotation_interval_days: 90,
            grace_period_days: 7,
            auto_rotate: false,
        }
    }
}

impl SshKeyManager {
    /// Create a new SSH key manager
    pub fn new(kms_backend: Arc<dyn KmsBackend>, rotation_policy: RotationPolicy) -> Self {
        Self {
            kms_backend,
            audit_log: Arc::new(RwLock::new(Vec::new())),
            rotation_policy,
        }
    }

    /// Store a new SSH key pair in KMS
    pub async fn store_ssh_key(
        &self,
        name: &str,
        private_key: &[u8],
        public_key: &[u8],
        purpose: SshKeyPurpose,
        tags: Vec<String>,
    ) -> Result<String, KmsError> {
        let key_id = Uuid::new_v4().to_string();
        let fingerprint = Self::calculate_fingerprint(public_key)?;

        // Store private key in KMS
        let private_key_data = KeyData {
            key_id: format!("{}_private", key_id),
            key_type: KeyType::Symmetric,  // Using symmetric for SSH keys
            key_material: private_key.to_vec(),
            metadata: KeyMetadata {
                name: format!("{}_private", name),
                created_at: Utc::now(),
                expires_at: None,
                tags: tags.clone(),
                owner: None,
                purpose: Some("ssh_private_key".to_string()),
            },
            encrypted: true,
        };

        self.kms_backend.store_key(private_key_data).await?;

        // Create SSH key pair record
        let ssh_key = SshKeyPair {
            key_id: key_id.clone(),
            name: name.to_string(),
            public_key: String::from_utf8_lossy(public_key).to_string(),
            private_key_id: format!("{}_private", key_id),
            fingerprint: fingerprint.clone(),
            created_at: Utc::now(),
            last_used: None,
            rotation_due: if self.rotation_policy.enabled {
                Some(Utc::now() + Duration::days(self.rotation_policy.rotation_interval_days as i64))
            } else {
                None
            },
            tags,
            purpose,
        };

        // Store SSH key metadata in KMS
        let metadata_key = KeyData {
            key_id: key_id.clone(),
            key_type: KeyType::Metadata,
            key_material: serde_json::to_vec(&ssh_key)
                .map_err(|e| KmsError::SerializationError(e.to_string()))?,
            metadata: KeyMetadata {
                name: name.to_string(),
                created_at: Utc::now(),
                expires_at: None,
                tags: vec!["ssh_key".to_string()],
                owner: None,
                purpose: Some("ssh_key_metadata".to_string()),
            },
            encrypted: false,
        };

        self.kms_backend.store_key(metadata_key).await?;

        // Audit log
        self.log_audit(SshKeyAuditEntry {
            timestamp: Utc::now(),
            key_id: key_id.clone(),
            action: SshKeyAction::Created,
            user: None,
            ip_address: None,
            success: true,
            error_message: None,
        }).await;

        Ok(key_id)
    }

    /// Retrieve SSH key pair from KMS
    pub async fn retrieve_ssh_key(&self, key_id: &str) -> Result<(Vec<u8>, Vec<u8>), KmsError> {
        // Get metadata
        let metadata_key = self.kms_backend.get_key(key_id).await?
            .ok_or_else(|| KmsError::KeyNotFound(key_id.to_string()))?;

        let ssh_key: SshKeyPair = serde_json::from_slice(&metadata_key.key_material)
            .map_err(|e| KmsError::SerializationError(e.to_string()))?;

        // Get private key
        let private_key_data = self.kms_backend.get_key(&ssh_key.private_key_id).await?
            .ok_or_else(|| KmsError::KeyNotFound(ssh_key.private_key_id.clone()))?;

        let public_key = ssh_key.public_key.as_bytes().to_vec();
        let private_key = private_key_data.key_material;

        // Update last used
        self.update_last_used(key_id).await?;

        // Audit log
        self.log_audit(SshKeyAuditEntry {
            timestamp: Utc::now(),
            key_id: key_id.to_string(),
            action: SshKeyAction::Retrieved,
            user: None,
            ip_address: None,
            success: true,
            error_message: None,
        }).await;

        Ok((private_key, public_key))
    }

    /// List all SSH keys
    pub async fn list_ssh_keys(&self) -> Result<Vec<SshKeyMetadata>, KmsError> {
        let all_keys = self.kms_backend.list_keys().await?;

        let mut ssh_keys = Vec::new();
        for key in all_keys {
            if key.metadata.purpose == Some("ssh_key_metadata".to_string()) {
                if let Ok(ssh_key) = serde_json::from_slice::<SshKeyPair>(&key.key_material) {
                    ssh_keys.push(SshKeyMetadata {
                        key_id: ssh_key.key_id,
                        name: ssh_key.name,
                        fingerprint: ssh_key.fingerprint,
                        created_at: ssh_key.created_at,
                        last_used: ssh_key.last_used,
                        rotation_due: ssh_key.rotation_due,
                        purpose: ssh_key.purpose,
                    });
                }
            }
        }

        Ok(ssh_keys)
    }

    /// Rotate SSH key
    pub async fn rotate_ssh_key(&self, key_id: &str) -> Result<String, KmsError> {
        // Get current key
        let current_key_data = self.kms_backend.get_key(key_id).await?
            .ok_or_else(|| KmsError::KeyNotFound(key_id.to_string()))?;

        let current_key: SshKeyPair = serde_json::from_slice(&current_key_data.key_material)
            .map_err(|e| KmsError::SerializationError(e.to_string()))?;

        // Generate new key pair (placeholder - would use actual SSH key generation)
        let (new_private, new_public) = self.generate_ssh_key_pair()?;

        // Store new key
        let new_key_id = self.store_ssh_key(
            &format!("{}_rotated", current_key.name),
            &new_private,
            &new_public,
            current_key.purpose,
            current_key.tags,
        ).await?;

        // Delete old key after grace period (in production, would schedule this)
        // For now, just mark as rotated in audit log

        // Audit log
        self.log_audit(SshKeyAuditEntry {
            timestamp: Utc::now(),
            key_id: key_id.to_string(),
            action: SshKeyAction::Rotated,
            user: None,
            ip_address: None,
            success: true,
            error_message: None,
        }).await;

        Ok(new_key_id)
    }

    /// Delete SSH key
    pub async fn delete_ssh_key(&self, key_id: &str) -> Result<(), KmsError> {
        // Get metadata to find private key ID
        let metadata_key = self.kms_backend.get_key(key_id).await?
            .ok_or_else(|| KmsError::KeyNotFound(key_id.to_string()))?;

        let ssh_key: SshKeyPair = serde_json::from_slice(&metadata_key.key_material)
            .map_err(|e| KmsError::SerializationError(e.to_string()))?;

        // Delete private key
        self.kms_backend.delete_key(&ssh_key.private_key_id).await?;

        // Delete metadata
        self.kms_backend.delete_key(key_id).await?;

        // Audit log
        self.log_audit(SshKeyAuditEntry {
            timestamp: Utc::now(),
            key_id: key_id.to_string(),
            action: SshKeyAction::Deleted,
            user: None,
            ip_address: None,
            success: true,
            error_message: None,
        }).await;

        Ok(())
    }

    /// Get audit log for a specific key
    pub async fn get_audit_log(&self, key_id: &str) -> Result<Vec<SshKeyAuditEntry>, KmsError> {
        let log = self.audit_log.read().await;
        Ok(log.iter()
            .filter(|entry| entry.key_id == key_id)
            .cloned()
            .collect())
    }

    /// Get all SSH keys due for rotation
    pub async fn get_keys_due_for_rotation(&self) -> Result<Vec<SshKeyMetadata>, KmsError> {
        let all_keys = self.list_ssh_keys().await?;
        let now = Utc::now();

        Ok(all_keys.into_iter()
            .filter(|key| {
                key.rotation_due
                    .map(|due_date| due_date <= now)
                    .unwrap_or(false)
            })
            .collect())
    }

    // Helper methods

    async fn update_last_used(&self, key_id: &str) -> Result<(), KmsError> {
        let metadata_key = self.kms_backend.get_key(key_id).await?
            .ok_or_else(|| KmsError::KeyNotFound(key_id.to_string()))?;

        let mut ssh_key: SshKeyPair = serde_json::from_slice(&metadata_key.key_material)
            .map_err(|e| KmsError::SerializationError(e.to_string()))?;

        ssh_key.last_used = Some(Utc::now());

        let updated_key = KeyData {
            key_material: serde_json::to_vec(&ssh_key)
                .map_err(|e| KmsError::SerializationError(e.to_string()))?,
            ..metadata_key
        };

        self.kms_backend.store_key(updated_key).await?;

        Ok(())
    }

    fn calculate_fingerprint(public_key: &[u8]) -> Result<String, KmsError> {
        use sha2::{Sha256, Digest};

        let mut hasher = Sha256::new();
        hasher.update(public_key);
        let result = hasher.finalize();

        Ok(format!("SHA256:{}", base64::encode(&result[..16])))
    }

    fn generate_ssh_key_pair(&self) -> Result<(Vec<u8>, Vec<u8>), KmsError> {
        // Placeholder - in production, would use proper SSH key generation
        // For now, return dummy keys for demonstration
        use rsa::{RsaPrivateKey, RsaPublicKey, pkcs8::{EncodePrivateKey, EncodePublicKey}};
        use rand::rngs::OsRng;

        let mut rng = OsRng;
        let bits = 2048;
        let private_key = RsaPrivateKey::new(&mut rng, bits)
            .map_err(|e| KmsError::CryptoError(e.to_string()))?;
        let public_key = RsaPublicKey::from(&private_key);

        let private_pem = private_key.to_pkcs8_pem(rsa::pkcs8::LineEnding::LF)
            .map_err(|e| KmsError::CryptoError(e.to_string()))?;
        let public_pem = public_key.to_public_key_pem(rsa::pkcs8::LineEnding::LF)
            .map_err(|e| KmsError::CryptoError(e.to_string()))?;

        Ok((private_pem.as_bytes().to_vec(), public_pem.as_bytes().to_vec()))
    }

    async fn log_audit(&self, entry: SshKeyAuditEntry) {
        let mut log = self.audit_log.write().await;
        log.push(entry);

        // Keep only last 10000 entries
        if log.len() > 10000 {
            log.drain(0..1000);
        }
    }
}

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

    // Tests would go here
}
