//! Hybrid KMS Backend
//!
//! Combines local and remote KMS backends with intelligent caching,
//! failover, and synchronization capabilities.

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

use async_trait::async_trait;
use chrono::{DateTime, Utc};
use tokio::sync::RwLock;
use uuid::Uuid;

use crate::kms::cache::CacheManager;
use crate::kms::error::KmsError;
use crate::kms::local::LocalKmsBackend;
use crate::kms::remote::RemoteKmsBackend;
use crate::kms::traits::{CacheBackend as CacheBackendTrait, KmsBackend};
use crate::kms::types::{HealthStatus, KeyData, KmsMode};
use crate::kms::KmsConfig;

/// Hybrid KMS backend that combines local and remote storage
pub struct HybridKmsBackend {
    local: LocalKmsBackend,
    remote: RemoteKmsBackend,
    cache: CacheManager,
    config: KmsConfig,
    sync_state: Arc<RwLock<SyncState>>,
    fallback_mode: Arc<RwLock<FallbackMode>>,
}

/// Synchronization state tracking
#[derive(Debug, Clone)]
struct SyncState {
    last_sync: Option<DateTime<Utc>>,
    pending_operations: Vec<PendingOperation>,
    sync_in_progress: bool,
    sync_failures: u32,
    last_sync_error: Option<String>,
}

/// Pending operations for offline synchronization
/// Timestamp used for operation ordering and retry logic
#[derive(Debug, Clone)]
#[allow(dead_code)]
struct PendingOperation {
    id: Uuid,
    operation_type: OperationType,
    key_id: String,
    data: Option<Vec<u8>>,
    timestamp: DateTime<Utc>,
    retries: u32,
}

/// Operation types for sync tracking
#[derive(Debug, Clone)]
enum OperationType {
    Store,
    Update,
    Delete,
}

/// Current fallback mode
#[derive(Debug, Clone, PartialEq)]
enum FallbackMode {
    /// Normal operation - both local and remote available
    Normal,
    /// Local only - remote unavailable
    LocalOnly,
    /// Remote only - local unavailable
    RemoteOnly,
    /// Emergency mode - significant issues detected
    Emergency,
}

impl HybridKmsBackend {
    /// Create a new hybrid KMS backend
    pub async fn new(config: &KmsConfig) -> Result<Self, KmsError> {
        // Validate hybrid mode configuration
        if config.mode != KmsMode::Hybrid {
            return Err(KmsError::config("Hybrid backend requires KmsMode::Hybrid"));
        }

        // Initialize backends
        let local = LocalKmsBackend::new(&config.local).await?;
        let remote = RemoteKmsBackend::new(&config.remote).await?;
        let cache = CacheManager::new(config.cache.clone()).await?;

        let sync_state = Arc::new(RwLock::new(SyncState {
            last_sync: None,
            pending_operations: Vec::new(),
            sync_in_progress: false,
            sync_failures: 0,
            last_sync_error: None,
        }));

        let fallback_mode = Arc::new(RwLock::new(FallbackMode::Normal));

        let backend = Self {
            local,
            remote,
            cache,
            config: config.clone(),
            sync_state,
            fallback_mode,
        };

        // Perform initial health check to determine fallback mode
        backend.update_fallback_mode().await;

        Ok(backend)
    }

    /// Update fallback mode based on backend health
    async fn update_fallback_mode(&self) {
        let local_health = self.local.health_check().await;
        let remote_health = self.remote.health_check().await;

        let new_mode = match (local_health, remote_health) {
            (Ok(local), Ok(remote)) => {
                if local.healthy && remote.healthy {
                    FallbackMode::Normal
                } else if local.healthy {
                    FallbackMode::LocalOnly
                } else if remote.healthy {
                    FallbackMode::RemoteOnly
                } else {
                    FallbackMode::Emergency
                }
            }
            (Ok(local), Err(_)) => {
                if local.healthy {
                    FallbackMode::LocalOnly
                } else {
                    FallbackMode::Emergency
                }
            }
            (Err(_), Ok(remote)) => {
                if remote.healthy {
                    FallbackMode::RemoteOnly
                } else {
                    FallbackMode::Emergency
                }
            }
            (Err(_), Err(_)) => FallbackMode::Emergency,
        };

        let mut fallback_mode = self.fallback_mode.write().await;
        if *fallback_mode != new_mode {
            tracing::info!(
                "Fallback mode changed from {:?} to {:?}",
                *fallback_mode,
                new_mode
            );
            *fallback_mode = new_mode;
        }
    }

    /// Add operation to pending sync queue
    async fn add_pending_operation(
        &self,
        op_type: OperationType,
        key_id: String,
        data: Option<Vec<u8>>,
    ) {
        let mut sync_state = self.sync_state.write().await;
        sync_state.pending_operations.push(PendingOperation {
            id: Uuid::new_v4(),
            operation_type: op_type,
            key_id,
            data,
            timestamp: Utc::now(),
            retries: 0,
        });
    }

    /// Handle a single pending operation
    async fn process_pending_operation(
        &self,
        mut op: PendingOperation,
    ) -> Result<(Option<Uuid>, Option<PendingOperation>), KmsError> {
        let key_id = op.key_id.clone();

        match &op.operation_type {
            OperationType::Store => {
                let Some(data) = &op.data else {
                    return Ok((None, Some(op)));
                };

                let Ok(key_data) = serde_json::from_slice::<KeyData>(data) else {
                    tracing::error!("Failed to deserialize key data for sync");
                    return Ok((None, Some(op)));
                };

                match self.remote.store_key(key_data).await {
                    Ok(_) => Ok((Some(op.id), None)),
                    Err(e) => {
                        op.retries += 1;
                        if op.retries < 3 {
                            tracing::warn!(
                                "Failed to sync store operation for key {}: {}",
                                key_id,
                                e
                            );
                            Ok((None, Some(op)))
                        } else {
                            Ok((None, None))
                        }
                    }
                }
            }
            OperationType::Update => {
                let Some(data) = &op.data else {
                    return Ok((None, Some(op)));
                };

                let Ok(key_data) = serde_json::from_slice::<KeyData>(data) else {
                    tracing::error!("Failed to deserialize key data for sync");
                    return Ok((None, Some(op)));
                };

                match self.remote.update_key(&op.key_id, key_data).await {
                    Ok(_) => Ok((Some(op.id), None)),
                    Err(e) => {
                        op.retries += 1;
                        if op.retries < 3 {
                            tracing::warn!(
                                "Failed to sync update operation for key {}: {}",
                                key_id,
                                e
                            );
                            Ok((None, Some(op)))
                        } else {
                            Ok((None, None))
                        }
                    }
                }
            }
            OperationType::Delete => match self.remote.delete_key(&op.key_id).await {
                Ok(_) => Ok((Some(op.id), None)),
                Err(e) => {
                    op.retries += 1;
                    if op.retries < 3 {
                        tracing::warn!("Failed to sync delete operation for key {}: {}", key_id, e);
                        Ok((None, Some(op)))
                    } else {
                        Ok((None, None))
                    }
                }
            },
        }
    }

    /// Synchronize pending operations with remote backend
    async fn sync_pending_operations(&self) -> Result<(), KmsError> {
        let mut sync_state = self.sync_state.write().await;

        if sync_state.sync_in_progress {
            return Ok(()); // Already syncing
        }

        sync_state.sync_in_progress = true;
        let pending_ops = std::mem::take(&mut sync_state.pending_operations);
        drop(sync_state);

        let mut successful_ops = Vec::new();
        let mut failed_ops = Vec::new();

        for op in pending_ops {
            let (success_id, failed_op) = self.process_pending_operation(op).await?;
            if let Some(id) = success_id {
                successful_ops.push(id);
            }
            if let Some(op) = failed_op {
                failed_ops.push(op);
            }
        }

        // Update sync state
        let mut sync_state = self.sync_state.write().await;
        sync_state.sync_in_progress = false;
        sync_state.last_sync = Some(Utc::now());
        sync_state.pending_operations.extend(failed_ops);

        if successful_ops.is_empty() && !sync_state.pending_operations.is_empty() {
            sync_state.sync_failures += 1;
            sync_state.last_sync_error = Some("All sync operations failed".to_string());
        } else {
            sync_state.sync_failures = 0;
            sync_state.last_sync_error = None;
        }

        tracing::info!(
            "Sync completed: {} successful, {} failed",
            successful_ops.len(),
            sync_state.pending_operations.len()
        );

        Ok(())
    }

    /// Get key from most appropriate source based on current mode
    async fn get_key_smart(&self, key_id: &str) -> Result<Option<KeyData>, KmsError> {
        let cache_key = format!("key:{}", key_id);

        // Try cache first
        if let Ok(Some(cached_data)) = self.cache.get(&cache_key).await {
            if let Ok(key_data) = serde_json::from_slice::<KeyData>(&cached_data) {
                return Ok(Some(key_data));
            }
        }

        let fallback_mode = self.fallback_mode.read().await.clone();

        match fallback_mode {
            FallbackMode::Normal => self.get_key_normal(&cache_key, key_id).await,
            FallbackMode::LocalOnly => self.local.get_key(key_id).await,
            FallbackMode::RemoteOnly => self.remote.get_key(key_id).await,
            FallbackMode::Emergency => {
                // Try both and return any successful result
                match self.local.get_key(key_id).await {
                    Ok(result) => Ok(result),
                    Err(_) => self.remote.get_key(key_id).await,
                }
            }
        }
    }

    async fn get_key_normal(
        &self,
        cache_key: &str,
        key_id: &str,
    ) -> Result<Option<KeyData>, KmsError> {
        // Try remote first for consistency, fallback to local
        match self.remote.get_key(key_id).await {
            Ok(Some(key_data)) => {
                // Cache the result
                if let Ok(serialized) = serde_json::to_vec(&key_data) {
                    let _ = self
                        .cache
                        .set(
                            cache_key,
                            &serialized,
                            self.config.cache.default_ttl_seconds,
                        )
                        .await;
                }
                Ok(Some(key_data))
            }
            Ok(None) => {
                // Try local backend
                match self.local.get_key(key_id).await {
                    Ok(Some(key_data)) => {
                        // Cache and possibly sync to remote
                        let Ok(serialized) = serde_json::to_vec(&key_data) else {
                            return Ok(Some(key_data));
                        };
                        let _ = self
                            .cache
                            .set(
                                cache_key,
                                &serialized,
                                self.config.cache.default_ttl_seconds,
                            )
                            .await;
                        Ok(Some(key_data))
                    }
                    Ok(None) => Ok(None),
                    Err(e) => Err(e),
                }
            }
            Err(_) => {
                // Remote failed, try local only
                self.local.get_key(key_id).await
            }
        }
    }

    /// Store key to appropriate backends based on current mode
    async fn store_key_smart(&self, key: KeyData) -> Result<String, KmsError> {
        let fallback_mode = self.fallback_mode.read().await.clone();
        let key_id = key.key_id.clone();

        // Cache the key
        if let Ok(serialized) = serde_json::to_vec(&key) {
            let cache_key = format!("key:{}", key_id);
            let _ = self
                .cache
                .set(
                    &cache_key,
                    &serialized,
                    self.config.cache.default_ttl_seconds,
                )
                .await;
        }

        match fallback_mode {
            FallbackMode::Normal => {
                // Store to both backends
                let local_result = self.local.store_key(key.clone()).await;
                let remote_result = self.remote.store_key(key.clone()).await;

                match (local_result, remote_result) {
                    (Ok(_), Ok(_)) => Ok(key_id),
                    (Ok(_), Err(_)) => {
                        // Remote failed, add to pending sync
                        if let Ok(serialized) = serde_json::to_vec(&key) {
                            self.add_pending_operation(
                                OperationType::Store,
                                key_id.clone(),
                                Some(serialized),
                            )
                            .await;
                        }
                        Ok(key_id)
                    }
                    (Err(local_err), Ok(_)) => {
                        tracing::warn!("Local storage failed: {}", local_err);
                        Ok(key_id)
                    }
                    (Err(local_err), Err(remote_err)) => Err(KmsError::storage(format!(
                        "Both backends failed: local={}, remote={}",
                        local_err, remote_err
                    ))),
                }
            }
            FallbackMode::LocalOnly => {
                let result = self.local.store_key(key.clone()).await;
                if result.is_ok() {
                    // Add to pending sync for when remote comes back online
                    if let Ok(serialized) = serde_json::to_vec(&key) {
                        self.add_pending_operation(
                            OperationType::Store,
                            key_id.clone(),
                            Some(serialized),
                        )
                        .await;
                    }
                }
                result
            }
            FallbackMode::RemoteOnly => self.remote.store_key(key).await,
            FallbackMode::Emergency => {
                // Try both, succeed if either works
                match self.local.store_key(key.clone()).await {
                    Ok(result) => Ok(result),
                    Err(_) => self.remote.store_key(key).await,
                }
            }
        }
    }

    /// Background task for periodic synchronization and health checks
    /// TODO: This requires wrapping HybridKmsBackend in Arc<> to work properly
    pub async fn start_background_tasks(&self) {
        tracing::warn!(
            "Background tasks not yet implemented - requires Arc wrapper for shared ownership"
        );
        // Implementation temporarily disabled until Arc wrapper is added
        // The background tasks would handle:
        // - Periodic synchronization of pending operations
        // - Health checks and fallback mode updates
        // - Cleanup of old cache entries
    }
}

// Note: This implementation assumes we can clone HybridKmsBackend
// In practice, you'd need to implement Clone or use Arc<HybridKmsBackend>
// for the background task spawning to work properly.

#[async_trait]
impl KmsBackend for HybridKmsBackend {
    async fn initialize(&mut self) -> Result<(), KmsError> {
        // Initialize both backends
        let local_init = self.local.initialize().await;
        let remote_init = self.remote.initialize().await;

        match (local_init, remote_init) {
            (Ok(()), Ok(())) => {
                self.start_background_tasks().await;
                Ok(())
            }
            (Ok(()), Err(e)) => {
                tracing::warn!(
                    "Remote backend initialization failed: {}, operating in local-only mode",
                    e
                );
                let mut fallback_mode = self.fallback_mode.write().await;
                *fallback_mode = FallbackMode::LocalOnly;
                Ok(())
            }
            (Err(e), Ok(())) => {
                tracing::warn!(
                    "Local backend initialization failed: {}, operating in remote-only mode",
                    e
                );
                let mut fallback_mode = self.fallback_mode.write().await;
                *fallback_mode = FallbackMode::RemoteOnly;
                Ok(())
            }
            (Err(local_err), Err(remote_err)) => Err(KmsError::backend(
                "hybrid",
                format!(
                    "Both backends failed to initialize: local={}, remote={}",
                    local_err, remote_err
                ),
            )),
        }
    }

    async fn store_key(&self, key: KeyData) -> Result<String, KmsError> {
        self.store_key_smart(key).await
    }

    async fn get_key(&self, key_id: &str) -> Result<Option<KeyData>, KmsError> {
        self.get_key_smart(key_id).await
    }

    async fn update_key(&self, key_id: &str, key: KeyData) -> Result<(), KmsError> {
        let fallback_mode = self.fallback_mode.read().await.clone();

        match fallback_mode {
            FallbackMode::Normal => {
                // Update both backends
                let local_result = self.local.update_key(key_id, key.clone()).await;
                let remote_result = self.remote.update_key(key_id, key.clone()).await;

                // Update cache
                if let Ok(serialized) = serde_json::to_vec(&key) {
                    let cache_key = format!("key:{}", key_id);
                    let _ = self
                        .cache
                        .set(
                            &cache_key,
                            &serialized,
                            self.config.cache.default_ttl_seconds,
                        )
                        .await;
                }

                match (local_result, remote_result) {
                    (Ok(()), Ok(())) => Ok(()),
                    (Ok(()), Err(_)) => {
                        // Remote failed, add to pending sync
                        if let Ok(serialized) = serde_json::to_vec(&key) {
                            self.add_pending_operation(
                                OperationType::Update,
                                key_id.to_string(),
                                Some(serialized),
                            )
                            .await;
                        }
                        Ok(())
                    }
                    (Err(_), Ok(())) => Ok(()),
                    (Err(local_err), Err(remote_err)) => Err(KmsError::storage(format!(
                        "Both backends failed: local={}, remote={}",
                        local_err, remote_err
                    ))),
                }
            }
            FallbackMode::LocalOnly => {
                let result = self.local.update_key(key_id, key.clone()).await;
                if result.is_ok() {
                    // Add to pending sync
                    if let Ok(serialized) = serde_json::to_vec(&key) {
                        self.add_pending_operation(
                            OperationType::Update,
                            key_id.to_string(),
                            Some(serialized),
                        )
                        .await;
                    }
                }
                result
            }
            FallbackMode::RemoteOnly => self.remote.update_key(key_id, key).await,
            FallbackMode::Emergency => {
                // Try both
                match self.local.update_key(key_id, key.clone()).await {
                    Ok(()) => Ok(()),
                    Err(_) => self.remote.update_key(key_id, key).await,
                }
            }
        }
    }

    async fn delete_key(&self, key_id: &str) -> Result<bool, KmsError> {
        let fallback_mode = self.fallback_mode.read().await.clone();

        // Remove from cache
        let cache_key = format!("key:{}", key_id);
        let _ = self.cache.delete(&cache_key).await;

        match fallback_mode {
            FallbackMode::Normal => {
                let local_result = self.local.delete_key(key_id).await;
                let remote_result = self.remote.delete_key(key_id).await;

                match (local_result, remote_result) {
                    (Ok(local_deleted), Ok(remote_deleted)) => Ok(local_deleted || remote_deleted),
                    (Ok(deleted), Err(_)) => {
                        // Remote failed, add to pending sync
                        self.add_pending_operation(OperationType::Delete, key_id.to_string(), None)
                            .await;
                        Ok(deleted)
                    }
                    (Err(_), Ok(deleted)) => Ok(deleted),
                    (Err(_), Err(_)) => Ok(false), // Neither backend worked, assume not found
                }
            }
            FallbackMode::LocalOnly => {
                let result = self.local.delete_key(key_id).await;
                if let Ok(deleted) = &result {
                    if *deleted {
                        self.add_pending_operation(OperationType::Delete, key_id.to_string(), None)
                            .await;
                    }
                }
                result
            }
            FallbackMode::RemoteOnly => self.remote.delete_key(key_id).await,
            FallbackMode::Emergency => match self.local.delete_key(key_id).await {
                Ok(deleted) => Ok(deleted),
                Err(_) => self.remote.delete_key(key_id).await,
            },
        }
    }

    async fn list_keys(
        &self,
        filters: Option<HashMap<String, String>>,
    ) -> Result<Vec<String>, KmsError> {
        let fallback_mode = self.fallback_mode.read().await.clone();

        match fallback_mode {
            FallbackMode::Normal => {
                // Merge results from both backends
                let local_keys = self
                    .local
                    .list_keys(filters.clone())
                    .await
                    .unwrap_or_default();
                let remote_keys = self.remote.list_keys(filters).await.unwrap_or_default();

                let mut all_keys = local_keys;
                for key in remote_keys {
                    if !all_keys.contains(&key) {
                        all_keys.push(key);
                    }
                }

                Ok(all_keys)
            }
            FallbackMode::LocalOnly => self.local.list_keys(filters).await,
            FallbackMode::RemoteOnly => self.remote.list_keys(filters).await,
            FallbackMode::Emergency => {
                // Try local first, then remote
                match self.local.list_keys(filters.clone()).await {
                    Ok(keys) => Ok(keys),
                    Err(_) => self.remote.list_keys(filters).await,
                }
            }
        }
    }

    async fn key_exists(&self, key_id: &str) -> Result<bool, KmsError> {
        self.get_key(key_id).await.map(|opt| opt.is_some())
    }

    async fn encrypt(
        &self,
        key_id: &str,
        plaintext: &[u8],
        context: Option<HashMap<String, String>>,
    ) -> Result<Vec<u8>, KmsError> {
        // For encryption, we prefer remote for consistency, but fallback to local
        let fallback_mode = self.fallback_mode.read().await.clone();

        match fallback_mode {
            FallbackMode::Normal | FallbackMode::RemoteOnly => {
                match self
                    .remote
                    .encrypt(key_id, plaintext, context.clone())
                    .await
                {
                    Ok(result) => Ok(result),
                    Err(_) if matches!(fallback_mode, FallbackMode::Normal) => {
                        self.local.encrypt(key_id, plaintext, context).await
                    }
                    Err(e) => Err(e),
                }
            }
            FallbackMode::LocalOnly => self.local.encrypt(key_id, plaintext, context).await,
            FallbackMode::Emergency => {
                // Try both
                match self
                    .remote
                    .encrypt(key_id, plaintext, context.clone())
                    .await
                {
                    Ok(result) => Ok(result),
                    Err(_) => self.local.encrypt(key_id, plaintext, context).await,
                }
            }
        }
    }

    async fn decrypt(
        &self,
        key_id: &str,
        ciphertext: &[u8],
        context: Option<HashMap<String, String>>,
    ) -> Result<Vec<u8>, KmsError> {
        // For decryption, try the same logic as encryption
        let fallback_mode = self.fallback_mode.read().await.clone();

        match fallback_mode {
            FallbackMode::Normal | FallbackMode::RemoteOnly => {
                match self
                    .remote
                    .decrypt(key_id, ciphertext, context.clone())
                    .await
                {
                    Ok(result) => Ok(result),
                    Err(_) if matches!(fallback_mode, FallbackMode::Normal) => {
                        self.local.decrypt(key_id, ciphertext, context).await
                    }
                    Err(e) => Err(e),
                }
            }
            FallbackMode::LocalOnly => self.local.decrypt(key_id, ciphertext, context).await,
            FallbackMode::Emergency => {
                // Try both
                match self
                    .remote
                    .decrypt(key_id, ciphertext, context.clone())
                    .await
                {
                    Ok(result) => Ok(result),
                    Err(_) => self.local.decrypt(key_id, ciphertext, context).await,
                }
            }
        }
    }

    async fn sign(
        &self,
        key_id: &str,
        data: &[u8],
        algorithm: Option<String>,
    ) -> Result<Vec<u8>, KmsError> {
        let fallback_mode = self.fallback_mode.read().await.clone();

        match fallback_mode {
            FallbackMode::Normal | FallbackMode::RemoteOnly => {
                match self.remote.sign(key_id, data, algorithm.clone()).await {
                    Ok(result) => Ok(result),
                    Err(_) if matches!(fallback_mode, FallbackMode::Normal) => {
                        self.local.sign(key_id, data, algorithm).await
                    }
                    Err(e) => Err(e),
                }
            }
            FallbackMode::LocalOnly => self.local.sign(key_id, data, algorithm).await,
            FallbackMode::Emergency => {
                match self.remote.sign(key_id, data, algorithm.clone()).await {
                    Ok(result) => Ok(result),
                    Err(_) => self.local.sign(key_id, data, algorithm).await,
                }
            }
        }
    }

    async fn verify(
        &self,
        key_id: &str,
        data: &[u8],
        signature: &[u8],
        algorithm: Option<String>,
    ) -> Result<bool, KmsError> {
        let fallback_mode = self.fallback_mode.read().await.clone();

        match fallback_mode {
            FallbackMode::Normal | FallbackMode::RemoteOnly => {
                match self
                    .remote
                    .verify(key_id, data, signature, algorithm.clone())
                    .await
                {
                    Ok(result) => Ok(result),
                    Err(_) if matches!(fallback_mode, FallbackMode::Normal) => {
                        self.local.verify(key_id, data, signature, algorithm).await
                    }
                    Err(e) => Err(e),
                }
            }
            FallbackMode::LocalOnly => self.local.verify(key_id, data, signature, algorithm).await,
            FallbackMode::Emergency => {
                match self
                    .remote
                    .verify(key_id, data, signature, algorithm.clone())
                    .await
                {
                    Ok(result) => Ok(result),
                    Err(_) => self.local.verify(key_id, data, signature, algorithm).await,
                }
            }
        }
    }

    async fn generate_data_key(
        &self,
        key_id: &str,
        key_spec: &str,
        context: Option<HashMap<String, String>>,
    ) -> Result<(Vec<u8>, Vec<u8>), KmsError> {
        let fallback_mode = self.fallback_mode.read().await.clone();

        match fallback_mode {
            FallbackMode::Normal | FallbackMode::RemoteOnly => {
                match self
                    .remote
                    .generate_data_key(key_id, key_spec, context.clone())
                    .await
                {
                    Ok(result) => Ok(result),
                    Err(_) if matches!(fallback_mode, FallbackMode::Normal) => {
                        self.local
                            .generate_data_key(key_id, key_spec, context)
                            .await
                    }
                    Err(e) => Err(e),
                }
            }
            FallbackMode::LocalOnly => {
                self.local
                    .generate_data_key(key_id, key_spec, context)
                    .await
            }
            FallbackMode::Emergency => {
                match self
                    .remote
                    .generate_data_key(key_id, key_spec, context.clone())
                    .await
                {
                    Ok(result) => Ok(result),
                    Err(_) => {
                        self.local
                            .generate_data_key(key_id, key_spec, context)
                            .await
                    }
                }
            }
        }
    }

    async fn health_check(&self) -> Result<HealthStatus, KmsError> {
        let local_health = self.local.health_check().await;
        let remote_health = self.remote.health_check().await;
        let cache_health = self.cache.health_check().await;

        let sync_state = self.sync_state.read().await;
        let fallback_mode = self.fallback_mode.read().await.clone();

        let mut metrics = HashMap::new();
        metrics.insert(
            "pending_sync_operations".to_string(),
            sync_state.pending_operations.len() as f64,
        );
        metrics.insert("sync_failures".to_string(), sync_state.sync_failures as f64);

        if let Some(last_sync) = sync_state.last_sync {
            metrics.insert(
                "seconds_since_last_sync".to_string(),
                (Utc::now() - last_sync).num_seconds() as f64,
            );
        }

        let overall_healthy = match (&local_health, &remote_health, &cache_health) {
            (Ok(l), Ok(r), Ok(_c)) => l.healthy || r.healthy, // At least one backend healthy
            (Ok(l), _, Ok(_c)) => l.healthy,                  // Local and cache work
            (_, Ok(r), Ok(_c)) => r.healthy,                  // Remote and cache work
            _ => false,                                       // Nothing works properly
        };

        let message = format!(
            "Hybrid KMS - Mode: {:?}, Local: {}, Remote: {}, Cache: {}, Pending sync: {}",
            fallback_mode,
            local_health.as_ref().map(|h| h.healthy).unwrap_or(false),
            remote_health.as_ref().map(|h| h.healthy).unwrap_or(false),
            cache_health.as_ref().map(|h| h.healthy).unwrap_or(false),
            sync_state.pending_operations.len()
        );

        Ok(HealthStatus {
            healthy: overall_healthy,
            message,
            last_check: Utc::now(),
            metrics,
        })
    }

    async fn shutdown(&mut self) -> Result<(), KmsError> {
        // Attempt final sync before shutdown
        let _ = self.sync_pending_operations().await;

        let local_shutdown = self.local.shutdown().await;
        let remote_shutdown = self.remote.shutdown().await;

        match (local_shutdown, remote_shutdown) {
            (Ok(()), Ok(())) => Ok(()),
            (Err(e), Ok(())) | (Ok(()), Err(e)) => {
                tracing::warn!("One backend failed to shutdown cleanly: {}", e);
                Ok(())
            }
            (Err(local_err), Err(remote_err)) => Err(KmsError::backend(
                "hybrid",
                format!(
                    "Both backends failed to shutdown: local={}, remote={}",
                    local_err, remote_err
                ),
            )),
        }
    }
}
