//! Local SQLite KMS Backend
//!
//! Implements a local KMS backend using SQLite with AES-256-GCM encryption
//! for secure key storage and management.

use std::collections::HashMap;

use aes_gcm::{AeadCore, AeadInPlace, Aes256Gcm, Key, KeyInit, Nonce};
use async_trait::async_trait;
use chrono::{DateTime, Utc};
use hkdf::Hkdf;
use sha2::Sha256;
use sqlx::{sqlite::SqlitePool, Pool, Row, Sqlite};
use zeroize::ZeroizeOnDrop;

use crate::kms::config::LocalConfig;
use crate::kms::error::KmsError;
use crate::kms::traits::KmsBackend;
use crate::kms::types::{
    HealthStatus, KeyAlgorithm, KeyData, KeyMetadata, KeyStatus, KeyType, KeyUsage, SecretBytes,
};

/// Local SQLite KMS backend implementation
pub struct LocalKmsBackend {
    pool: Pool<Sqlite>,
    encryption_key: EncryptionKey,
    /// Configuration used for initialization and future extensions
    #[allow(dead_code)]
    config: LocalConfig,
}

/// Master encryption key for database encryption
#[derive(ZeroizeOnDrop)]
struct EncryptionKey {
    #[zeroize(skip)]
    key: Key<Aes256Gcm>,
}

impl EncryptionKey {
    /// Derive encryption key from master key material
    fn derive_from_master(master_key: &[u8], salt: &[u8]) -> Result<Self, KmsError> {
        let hkdf = Hkdf::<Sha256>::new(Some(salt), master_key);
        let mut key_bytes = [0u8; 32]; // 256 bits
        hkdf.expand(b"KMS-LOCAL-STORAGE-KEY", &mut key_bytes)
            .map_err(|e| KmsError::crypto("HKDF", format!("Key derivation failed: {}", e)))?;

        let key = Key::<Aes256Gcm>::from_slice(&key_bytes);
        Ok(Self { key: *key })
    }

    /// Encrypt data using AES-256-GCM
    fn encrypt(&self, plaintext: &[u8], aad: &[u8]) -> Result<Vec<u8>, KmsError> {
        use rand::RngCore;
        let cipher = Aes256Gcm::new(&self.key);
        let mut nonce_bytes = [0u8; 12];
        rand::rng().fill_bytes(&mut nonce_bytes);
        let nonce = Nonce::from_slice(&nonce_bytes);

        let mut buffer = plaintext.to_vec();
        let tag = cipher
            .encrypt_in_place_detached(nonce, aad, &mut buffer)
            .map_err(|e| KmsError::crypto("AES-GCM", format!("Encryption failed: {}", e)))?;

        // Format: nonce (12 bytes) + tag (16 bytes) + ciphertext
        let mut result = Vec::with_capacity(12 + 16 + buffer.len());
        result.extend_from_slice(nonce);
        result.extend_from_slice(&tag);
        result.extend_from_slice(&buffer);

        Ok(result)
    }

    /// Decrypt data using AES-256-GCM
    fn decrypt(&self, ciphertext: &[u8], aad: &[u8]) -> Result<Vec<u8>, KmsError> {
        if ciphertext.len() < 28 {
            // nonce (12) + tag (16) = 28 minimum
            return Err(KmsError::crypto("AES-GCM", "Ciphertext too short"));
        }

        let cipher = Aes256Gcm::new(&self.key);
        let nonce = Nonce::from_slice(&ciphertext[0..12]);
        let tag = &ciphertext[12..28];
        let mut buffer = ciphertext[28..].to_vec();

        cipher
            .decrypt_in_place_detached(nonce, aad, &mut buffer, tag.into())
            .map_err(|e| KmsError::crypto("AES-GCM", format!("Decryption failed: {}", e)))?;

        Ok(buffer)
    }
}

impl LocalKmsBackend {
    /// Create a new local KMS backend
    pub async fn new(config: &LocalConfig) -> Result<Self, KmsError> {
        // Initialize database connection pool
        let pool = SqlitePool::connect(&format!("sqlite:{}", config.database_path))
            .await
            .map_err(|e| KmsError::storage(format!("Database connection failed: {}", e)))?;

        // Derive encryption key
        let master_key = Self::get_or_generate_master_key(config).await?;
        let salt = Self::get_or_generate_salt(config).await?;
        let encryption_key = EncryptionKey::derive_from_master(&master_key, &salt)?;

        let backend = Self {
            pool,
            encryption_key,
            config: config.clone(),
        };

        Ok(backend)
    }

    /// Initialize the database schema
    async fn init_schema(&self) -> Result<(), KmsError> {
        // Create keys table
        sqlx::query(
            r#"
            CREATE TABLE IF NOT EXISTS keys (
                key_id TEXT PRIMARY KEY,
                key_type TEXT NOT NULL,
                algorithm TEXT NOT NULL,
                key_size INTEGER NOT NULL,
                encrypted_key_material BLOB NOT NULL,
                usage_encrypt BOOLEAN NOT NULL DEFAULT 1,
                usage_decrypt BOOLEAN NOT NULL DEFAULT 1,
                usage_sign BOOLEAN NOT NULL DEFAULT 0,
                usage_verify BOOLEAN NOT NULL DEFAULT 0,
                usage_wrap BOOLEAN NOT NULL DEFAULT 0,
                usage_unwrap BOOLEAN NOT NULL DEFAULT 0,
                usage_derive BOOLEAN NOT NULL DEFAULT 0,
                usage_export BOOLEAN NOT NULL DEFAULT 0,
                metadata_name TEXT,
                metadata_description TEXT,
                metadata_owner TEXT,
                metadata_application TEXT,
                metadata_environment TEXT,
                metadata_version INTEGER NOT NULL DEFAULT 1,
                metadata_custom TEXT, -- JSON
                created_at TEXT NOT NULL,
                last_accessed TEXT,
                expires_at TEXT,
                status TEXT NOT NULL DEFAULT 'Active',
                tags TEXT, -- JSON
                UNIQUE(key_id)
            )
            "#,
        )
        .execute(&self.pool)
        .await
        .map_err(|e| KmsError::storage(format!("Schema creation failed: {}", e)))?;

        // Create indexes for performance
        sqlx::query("CREATE INDEX IF NOT EXISTS idx_keys_status ON keys(status)")
            .execute(&self.pool)
            .await
            .map_err(|e| KmsError::storage(format!("Index creation failed: {}", e)))?;

        sqlx::query("CREATE INDEX IF NOT EXISTS idx_keys_created_at ON keys(created_at)")
            .execute(&self.pool)
            .await
            .map_err(|e| KmsError::storage(format!("Index creation failed: {}", e)))?;

        sqlx::query("CREATE INDEX IF NOT EXISTS idx_keys_expires_at ON keys(expires_at)")
            .execute(&self.pool)
            .await
            .map_err(|e| KmsError::storage(format!("Index creation failed: {}", e)))?;

        Ok(())
    }

    /// Get or generate master key based on configuration
    async fn get_or_generate_master_key(config: &LocalConfig) -> Result<Vec<u8>, KmsError> {
        match config.master_key.source {
            crate::kms::KeySource::File => {
                let path = config
                    .master_key
                    .key_file_path
                    .as_ref()
                    .ok_or_else(|| KmsError::config("Master key file path not specified"))?;

                tokio::fs::read(path)
                    .await
                    .map_err(|e| KmsError::io(format!("Failed to read master key file: {}", e)))
            }
            crate::kms::KeySource::Environment => {
                let var_name = config.master_key.env_var_name.as_ref().ok_or_else(|| {
                    KmsError::config("Master key environment variable not specified")
                })?;

                let value = std::env::var(var_name).map_err(|e| {
                    KmsError::config(format!("Master key environment variable not found: {}", e))
                })?;

                Ok(value.into_bytes())
            }
            crate::kms::KeySource::Generated => {
                // Generate a secure random key
                use rand::RngCore;
                let mut key = vec![0u8; 32];
                rand::rng().fill_bytes(&mut key);

                // Generated keys are returned but not persisted by default.
                // Callers are responsible for saving the key if needed.
                Ok(key)
            }
            crate::kms::KeySource::Hsm => {
                // HSM key retrieval requires HSM module support
                Err(KmsError::config(
                    "HSM master key source not yet implemented",
                ))
            }
        }
    }

    /// Get or generate salt for key derivation
    async fn get_or_generate_salt(config: &LocalConfig) -> Result<Vec<u8>, KmsError> {
        if let Some(salt_str) = &config.master_key.salt {
            Ok(salt_str.as_bytes().to_vec())
        } else {
            // Generate a random salt
            use rand::RngCore;
            let mut salt = vec![0u8; 16];
            rand::rng().fill_bytes(&mut salt);
            Ok(salt)
        }
    }

    /// Encrypt key material for storage
    fn encrypt_key_material(&self, key_data: &KeyData) -> Result<Vec<u8>, KmsError> {
        let aad = format!("key:{}", key_data.key_id);
        self.encryption_key
            .encrypt(key_data.key_material.as_slice(), aad.as_bytes())
    }

    /// Decrypt key material from storage
    fn decrypt_key_material(
        &self,
        key_id: &str,
        encrypted_data: &[u8],
    ) -> Result<SecretBytes, KmsError> {
        let aad = format!("key:{}", key_id);
        let decrypted = self
            .encryption_key
            .decrypt(encrypted_data, aad.as_bytes())?;
        Ok(SecretBytes::new(decrypted))
    }

    /// Convert database row to KeyData
    fn row_to_key_data(&self, row: &sqlx::sqlite::SqliteRow) -> Result<KeyData, KmsError> {
        let key_id: String = row.get("key_id");
        let encrypted_material: Vec<u8> = row.get("encrypted_key_material");
        let key_material = self.decrypt_key_material(&key_id, &encrypted_material)?;

        // Parse key type
        let key_type_str: String = row.get("key_type");
        let key_type = match key_type_str.as_str() {
            "Symmetric" => KeyType::Symmetric,
            "Asymmetric" => KeyType::Asymmetric,
            "Hmac" => KeyType::Hmac,
            "Kdf" => KeyType::Kdf,
            "Dek" => KeyType::Dek,
            "Kek" => KeyType::Kek,
            _ => {
                return Err(KmsError::storage(format!(
                    "Invalid key type: {}",
                    key_type_str
                )))
            }
        };

        // Parse algorithm
        let algorithm_str: String = row.get("algorithm");
        let algorithm = match algorithm_str.as_str() {
            "AES-256-GCM" => KeyAlgorithm::Aes256Gcm,
            "AES-128-GCM" => KeyAlgorithm::Aes128Gcm,
            "ChaCha20Poly1305" => KeyAlgorithm::ChaCha20Poly1305,
            "RSA-2048" => KeyAlgorithm::Rsa2048,
            "RSA-4096" => KeyAlgorithm::Rsa4096,
            "ECDSA-P256" => KeyAlgorithm::EcdsaP256,
            "ECDSA-P384" => KeyAlgorithm::EcdsaP384,
            "Ed25519" => KeyAlgorithm::Ed25519,
            "HMAC-SHA256" => KeyAlgorithm::HmacSha256,
            "HMAC-SHA512" => KeyAlgorithm::HmacSha512,
            "HKDF-SHA256" => KeyAlgorithm::HkdfSha256,
            _ => {
                return Err(KmsError::storage(format!(
                    "Invalid algorithm: {}",
                    algorithm_str
                )))
            }
        };

        // Parse usage
        let usage = KeyUsage {
            encrypt: row.get("usage_encrypt"),
            decrypt: row.get("usage_decrypt"),
            sign: row.get("usage_sign"),
            verify: row.get("usage_verify"),
            wrap: row.get("usage_wrap"),
            unwrap: row.get("usage_unwrap"),
            derive: row.get("usage_derive"),
            export: row.get("usage_export"),
        };

        // Parse metadata
        let custom_str: Option<String> = row.get("metadata_custom");
        let custom: HashMap<String, String> = if let Some(json_str) = custom_str {
            serde_json::from_str(&json_str).unwrap_or_default()
        } else {
            HashMap::new()
        };

        let metadata = KeyMetadata {
            name: row.get("metadata_name"),
            description: row.get("metadata_description"),
            owner: row.get("metadata_owner"),
            application: row.get("metadata_application"),
            environment: row.get("metadata_environment"),
            version: row.get("metadata_version"),
            custom,
        };

        // Parse timestamps
        let created_at_str: String = row.get("created_at");
        let created_at = DateTime::parse_from_rfc3339(&created_at_str)
            .map_err(|e| KmsError::storage(format!("Invalid created_at timestamp: {}", e)))?
            .with_timezone(&Utc);

        let last_accessed =
            if let Some(last_accessed_str) = row.get::<Option<String>, _>("last_accessed") {
                Some(
                    DateTime::parse_from_rfc3339(&last_accessed_str)
                        .map_err(|e| {
                            KmsError::storage(format!("Invalid last_accessed timestamp: {}", e))
                        })?
                        .with_timezone(&Utc),
                )
            } else {
                None
            };

        let expires_at = if let Some(expires_at_str) = row.get::<Option<String>, _>("expires_at") {
            Some(
                DateTime::parse_from_rfc3339(&expires_at_str)
                    .map_err(|e| KmsError::storage(format!("Invalid expires_at timestamp: {}", e)))?
                    .with_timezone(&Utc),
            )
        } else {
            None
        };

        // Parse status
        let status_str: String = row.get("status");
        let status = match status_str.as_str() {
            "Active" => KeyStatus::Active,
            "Pending" => KeyStatus::Pending,
            "Deactivated" => KeyStatus::Deactivated,
            "Compromised" => KeyStatus::Compromised,
            "Destroyed" => KeyStatus::Destroyed,
            "Archived" => KeyStatus::Archived,
            _ => {
                return Err(KmsError::storage(format!(
                    "Invalid key status: {}",
                    status_str
                )))
            }
        };

        // Parse tags
        let tags_str: Option<String> = row.get("tags");
        let tags: HashMap<String, String> = if let Some(json_str) = tags_str {
            serde_json::from_str(&json_str).unwrap_or_default()
        } else {
            HashMap::new()
        };

        Ok(KeyData {
            key_id,
            key_type,
            algorithm,
            usage,
            key_size: row.get("key_size"),
            key_material,
            metadata,
            created_at,
            last_accessed,
            expires_at,
            status,
            tags,
        })
    }
}

#[async_trait]
impl KmsBackend for LocalKmsBackend {
    async fn initialize(&mut self) -> Result<(), KmsError> {
        self.init_schema().await
    }

    async fn store_key(&self, key: KeyData) -> Result<String, KmsError> {
        // Check if key already exists
        if self.key_exists(&key.key_id).await? {
            return Err(KmsError::key_exists(&key.key_id));
        }

        // Encrypt key material
        let encrypted_material = self.encrypt_key_material(&key)?;

        // Serialize complex fields
        let custom_json = serde_json::to_string(&key.metadata.custom).map_err(|e| {
            KmsError::serialization(format!("Failed to serialize custom metadata: {}", e))
        })?;

        let tags_json = serde_json::to_string(&key.tags)
            .map_err(|e| KmsError::serialization(format!("Failed to serialize tags: {}", e)))?;

        // Insert into database
        sqlx::query(
            r#"
            INSERT INTO keys (
                key_id, key_type, algorithm, key_size, encrypted_key_material,
                usage_encrypt, usage_decrypt, usage_sign, usage_verify,
                usage_wrap, usage_unwrap, usage_derive, usage_export,
                metadata_name, metadata_description, metadata_owner,
                metadata_application, metadata_environment, metadata_version,
                metadata_custom, created_at, last_accessed, expires_at,
                status, tags
            ) VALUES (
                ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?
            )
            "#,
        )
        .bind(&key.key_id)
        .bind(format!("{:?}", key.key_type))
        .bind(format!("{:?}", key.algorithm))
        .bind(key.key_size as i32)
        .bind(&encrypted_material)
        .bind(key.usage.encrypt)
        .bind(key.usage.decrypt)
        .bind(key.usage.sign)
        .bind(key.usage.verify)
        .bind(key.usage.wrap)
        .bind(key.usage.unwrap)
        .bind(key.usage.derive)
        .bind(key.usage.export)
        .bind(&key.metadata.name)
        .bind(&key.metadata.description)
        .bind(&key.metadata.owner)
        .bind(&key.metadata.application)
        .bind(&key.metadata.environment)
        .bind(key.metadata.version as i32)
        .bind(&custom_json)
        .bind(key.created_at.to_rfc3339())
        .bind(key.last_accessed.map(|t| t.to_rfc3339()))
        .bind(key.expires_at.map(|t| t.to_rfc3339()))
        .bind(format!("{:?}", key.status))
        .bind(&tags_json)
        .execute(&self.pool)
        .await
        .map_err(|e| KmsError::storage(format!("Failed to store key: {}", e)))?;

        Ok(key.key_id)
    }

    async fn get_key(&self, key_id: &str) -> Result<Option<KeyData>, KmsError> {
        let row = sqlx::query("SELECT * FROM keys WHERE key_id = ?")
            .bind(key_id)
            .fetch_optional(&self.pool)
            .await
            .map_err(|e| KmsError::storage(format!("Failed to retrieve key: {}", e)))?;

        if let Some(row) = row {
            // Update last_accessed timestamp
            let now = Utc::now();
            sqlx::query("UPDATE keys SET last_accessed = ? WHERE key_id = ?")
                .bind(now.to_rfc3339())
                .bind(key_id)
                .execute(&self.pool)
                .await
                .map_err(|e| KmsError::storage(format!("Failed to update last_accessed: {}", e)))?;

            let mut key_data = self.row_to_key_data(&row)?;
            key_data.last_accessed = Some(now);
            Ok(Some(key_data))
        } else {
            Ok(None)
        }
    }

    async fn update_key(&self, key_id: &str, key: KeyData) -> Result<(), KmsError> {
        if !self.key_exists(key_id).await? {
            return Err(KmsError::key_not_found(key_id));
        }

        // Encrypt key material
        let encrypted_material = self.encrypt_key_material(&key)?;

        // Serialize complex fields
        let custom_json = serde_json::to_string(&key.metadata.custom).map_err(|e| {
            KmsError::serialization(format!("Failed to serialize custom metadata: {}", e))
        })?;

        let tags_json = serde_json::to_string(&key.tags)
            .map_err(|e| KmsError::serialization(format!("Failed to serialize tags: {}", e)))?;

        // Update in database
        sqlx::query(
            r#"
            UPDATE keys SET
                key_type = ?, algorithm = ?, key_size = ?, encrypted_key_material = ?,
                usage_encrypt = ?, usage_decrypt = ?, usage_sign = ?, usage_verify = ?,
                usage_wrap = ?, usage_unwrap = ?, usage_derive = ?, usage_export = ?,
                metadata_name = ?, metadata_description = ?, metadata_owner = ?,
                metadata_application = ?, metadata_environment = ?, metadata_version = ?,
                metadata_custom = ?, expires_at = ?, status = ?, tags = ?
            WHERE key_id = ?
            "#,
        )
        .bind(format!("{:?}", key.key_type))
        .bind(format!("{:?}", key.algorithm))
        .bind(key.key_size as i32)
        .bind(&encrypted_material)
        .bind(key.usage.encrypt)
        .bind(key.usage.decrypt)
        .bind(key.usage.sign)
        .bind(key.usage.verify)
        .bind(key.usage.wrap)
        .bind(key.usage.unwrap)
        .bind(key.usage.derive)
        .bind(key.usage.export)
        .bind(&key.metadata.name)
        .bind(&key.metadata.description)
        .bind(&key.metadata.owner)
        .bind(&key.metadata.application)
        .bind(&key.metadata.environment)
        .bind(key.metadata.version as i32)
        .bind(&custom_json)
        .bind(key.expires_at.map(|t| t.to_rfc3339()))
        .bind(format!("{:?}", key.status))
        .bind(&tags_json)
        .bind(key_id)
        .execute(&self.pool)
        .await
        .map_err(|e| KmsError::storage(format!("Failed to update key: {}", e)))?;

        Ok(())
    }

    async fn delete_key(&self, key_id: &str) -> Result<bool, KmsError> {
        let result = sqlx::query("DELETE FROM keys WHERE key_id = ?")
            .bind(key_id)
            .execute(&self.pool)
            .await
            .map_err(|e| KmsError::storage(format!("Failed to delete key: {}", e)))?;

        Ok(result.rows_affected() > 0)
    }

    async fn list_keys(
        &self,
        filters: Option<HashMap<String, String>>,
    ) -> Result<Vec<String>, KmsError> {
        let mut query = "SELECT key_id FROM keys WHERE 1=1".to_string();
        let mut bind_params = Vec::new();

        if let Some(filters) = filters {
            for (key, value) in filters {
                match key.as_str() {
                    "status" => {
                        query.push_str(" AND status = ?");
                        bind_params.push(value);
                    }
                    "key_type" => {
                        query.push_str(" AND key_type = ?");
                        bind_params.push(value);
                    }
                    "algorithm" => {
                        query.push_str(" AND algorithm = ?");
                        bind_params.push(value);
                    }
                    "owner" => {
                        query.push_str(" AND metadata_owner = ?");
                        bind_params.push(value);
                    }
                    "application" => {
                        query.push_str(" AND metadata_application = ?");
                        bind_params.push(value);
                    }
                    "environment" => {
                        query.push_str(" AND metadata_environment = ?");
                        bind_params.push(value);
                    }
                    _ => {} // Ignore unknown filters
                }
            }
        }

        query.push_str(" ORDER BY created_at DESC");

        let mut sql_query = sqlx::query(&query);
        for param in bind_params {
            sql_query = sql_query.bind(param);
        }

        let rows = sql_query
            .fetch_all(&self.pool)
            .await
            .map_err(|e| KmsError::storage(format!("Failed to list keys: {}", e)))?;

        let key_ids = rows
            .iter()
            .map(|row| row.get::<String, _>("key_id"))
            .collect();

        Ok(key_ids)
    }

    async fn key_exists(&self, key_id: &str) -> Result<bool, KmsError> {
        let row = sqlx::query("SELECT COUNT(*) as count FROM keys WHERE key_id = ?")
            .bind(key_id)
            .fetch_one(&self.pool)
            .await
            .map_err(|e| KmsError::storage(format!("Failed to check key existence: {}", e)))?;

        let count: i32 = row.get("count");
        Ok(count > 0)
    }

    async fn encrypt(
        &self,
        key_id: &str,
        plaintext: &[u8],
        context: Option<HashMap<String, String>>,
    ) -> Result<Vec<u8>, KmsError> {
        let key_data = self
            .get_key(key_id)
            .await?
            .ok_or_else(|| KmsError::key_not_found(key_id))?;

        if !key_data.usage.encrypt {
            return Err(KmsError::permission_denied("encrypt", key_id));
        }

        if key_data.status != KeyStatus::Active {
            return Err(KmsError::invalid_key(key_id, "Key is not active"));
        }

        // Check for key expiration
        if let Some(expires_at) = key_data.expires_at {
            if Utc::now() > expires_at {
                return Err(KmsError::key_expired(key_id));
            }
        }

        match key_data.algorithm {
            KeyAlgorithm::Aes256Gcm | KeyAlgorithm::Aes128Gcm => {
                use rand::RngCore;
                // Use the key material directly for AES-GCM
                let key = Key::<Aes256Gcm>::from_slice(key_data.key_material.as_slice());
                let cipher = Aes256Gcm::new(key);
                let mut nonce_bytes = [0u8; 12];
                rand::rng().fill_bytes(&mut nonce_bytes);
                let nonce = Nonce::from_slice(&nonce_bytes);

                // Build AAD from context
                let aad = if let Some(ctx) = context {
                    serde_json::to_vec(&ctx).map_err(|e| {
                        KmsError::serialization(format!("Failed to serialize context: {}", e))
                    })?
                } else {
                    Vec::new()
                };

                let mut buffer = plaintext.to_vec();
                let tag = cipher
                    .encrypt_in_place_detached(nonce, &aad, &mut buffer)
                    .map_err(|e| {
                        KmsError::crypto("AES-GCM", format!("Encryption failed: {}", e))
                    })?;

                // Format: nonce + tag + ciphertext
                let mut result = Vec::with_capacity(12 + 16 + buffer.len());
                result.extend_from_slice(nonce);
                result.extend_from_slice(&tag);
                result.extend_from_slice(&buffer);

                Ok(result)
            }
            _ => Err(KmsError::crypto(
                "encrypt",
                format!(
                    "Algorithm {:?} not supported for encryption",
                    key_data.algorithm
                ),
            )),
        }
    }

    async fn decrypt(
        &self,
        key_id: &str,
        ciphertext: &[u8],
        context: Option<HashMap<String, String>>,
    ) -> Result<Vec<u8>, KmsError> {
        let key_data = self
            .get_key(key_id)
            .await?
            .ok_or_else(|| KmsError::key_not_found(key_id))?;

        if !key_data.usage.decrypt {
            return Err(KmsError::permission_denied("decrypt", key_id));
        }

        if key_data.status != KeyStatus::Active {
            return Err(KmsError::invalid_key(key_id, "Key is not active"));
        }

        // Check for key expiration
        if let Some(expires_at) = key_data.expires_at {
            if Utc::now() > expires_at {
                return Err(KmsError::key_expired(key_id));
            }
        }

        match key_data.algorithm {
            KeyAlgorithm::Aes256Gcm | KeyAlgorithm::Aes128Gcm => {
                if ciphertext.len() < 28 {
                    return Err(KmsError::crypto("AES-GCM", "Ciphertext too short"));
                }

                let key = Key::<Aes256Gcm>::from_slice(key_data.key_material.as_slice());
                let cipher = Aes256Gcm::new(key);
                let nonce = Nonce::from_slice(&ciphertext[0..12]);
                let tag = &ciphertext[12..28];
                let mut buffer = ciphertext[28..].to_vec();

                // Build AAD from context
                let aad = if let Some(ctx) = context {
                    serde_json::to_vec(&ctx).map_err(|e| {
                        KmsError::serialization(format!("Failed to serialize context: {}", e))
                    })?
                } else {
                    Vec::new()
                };

                cipher
                    .decrypt_in_place_detached(nonce, &aad, &mut buffer, tag.into())
                    .map_err(|e| {
                        KmsError::crypto("AES-GCM", format!("Decryption failed: {}", e))
                    })?;

                Ok(buffer)
            }
            _ => Err(KmsError::crypto(
                "decrypt",
                format!(
                    "Algorithm {:?} not supported for decryption",
                    key_data.algorithm
                ),
            )),
        }
    }

    async fn sign(
        &self,
        _key_id: &str,
        _data: &[u8],
        _algorithm: Option<String>,
    ) -> Result<Vec<u8>, KmsError> {
        // TODO: Implement signing for asymmetric keys
        Err(KmsError::crypto("sign", "Signing not yet implemented"))
    }

    async fn verify(
        &self,
        _key_id: &str,
        _data: &[u8],
        _signature: &[u8],
        _algorithm: Option<String>,
    ) -> Result<bool, KmsError> {
        // TODO: Implement signature verification
        Err(KmsError::crypto(
            "verify",
            "Signature verification not yet implemented",
        ))
    }

    async fn generate_data_key(
        &self,
        key_id: &str,
        key_spec: &str,
        context: Option<HashMap<String, String>>,
    ) -> Result<(Vec<u8>, Vec<u8>), KmsError> {
        let key_data = self
            .get_key(key_id)
            .await?
            .ok_or_else(|| KmsError::key_not_found(key_id))?;

        if key_data.key_type != KeyType::Kek {
            return Err(KmsError::invalid_key(
                key_id,
                "Key is not a key encryption key",
            ));
        }

        // Parse key spec (e.g., "AES_256", "AES_128")
        let key_length = match key_spec {
            "AES_256" => 32,
            "AES_128" => 16,
            _ => {
                return Err(KmsError::validation(
                    "key_spec",
                    "Unsupported key specification",
                ))
            }
        };

        // Generate random data key
        use rand::RngCore;
        let mut data_key = vec![0u8; key_length];
        rand::rng().fill_bytes(&mut data_key);

        // Encrypt data key using the KEK
        let encrypted_data_key = self.encrypt(key_id, &data_key, context).await?;

        Ok((data_key, encrypted_data_key))
    }

    async fn health_check(&self) -> Result<HealthStatus, KmsError> {
        // Try to execute a simple query
        let start = std::time::Instant::now();
        let result = sqlx::query("SELECT COUNT(*) as count FROM keys")
            .fetch_one(&self.pool)
            .await;

        let elapsed = start.elapsed();

        match result {
            Ok(row) => {
                let count: i32 = row.get("count");
                let mut metrics = HashMap::new();
                metrics.insert("query_time_ms".to_string(), elapsed.as_millis() as f64);
                metrics.insert("key_count".to_string(), count as f64);

                Ok(HealthStatus {
                    healthy: true,
                    message: format!("Local KMS backend healthy - {} keys stored", count),
                    last_check: Utc::now(),
                    metrics,
                })
            }
            Err(e) => Ok(HealthStatus::unhealthy(format!(
                "Database query failed: {}",
                e
            ))),
        }
    }

    async fn shutdown(&mut self) -> Result<(), KmsError> {
        self.pool.close().await;
        Ok(())
    }
}
