//! State management and persistence

use crate::{EtlError, Result};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::RwLock;

/// State management configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StateConfig {
    pub backend: StateBackendType,
    pub checkpoint_interval_ms: u64,
    pub cleanup_interval_ms: u64,
}

impl Default for StateConfig {
    fn default() -> Self {
        Self {
            backend: StateBackendType::InMemory,
            checkpoint_interval_ms: 30000,
            cleanup_interval_ms: 300000,
        }
    }
}

/// State backend types
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum StateBackendType {
    InMemory,
    Redis { connection_string: String },
    PostgreSQL { connection_string: String },
}

/// State manager  
#[derive(Debug)]
pub struct StateManager {
    config: StateConfig,
    backend: Arc<dyn StateBackend>,
    state_cache: Arc<RwLock<HashMap<String, StateEntry>>>,
}

/// State entry
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StateEntry {
    pub key: String,
    pub value: serde_json::Value,
    pub timestamp: DateTime<Utc>,
    pub ttl: Option<u64>,
}

/// State backend trait  
pub trait StateBackend: Send + Sync + std::fmt::Debug {
    fn get(
        &self,
        key: &str,
    ) -> std::pin::Pin<
        Box<dyn std::future::Future<Output = Result<Option<serde_json::Value>>> + Send + '_>,
    >;
    fn set(
        &self,
        key: &str,
        value: serde_json::Value,
    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send + '_>>;
    fn delete(
        &self,
        key: &str,
    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send + '_>>;
    fn list_keys(
        &self,
        prefix: &str,
    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<Vec<String>>> + Send + '_>>;
}

/// State snapshot
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StateSnapshot {
    pub snapshot_id: String,
    pub timestamp: DateTime<Utc>,
    pub entries: HashMap<String, serde_json::Value>,
}

/// State recovery
#[derive(Debug)]
pub struct StateRecovery {
    state_manager: Arc<StateManager>,
}

/// Redis state manager
#[derive(Debug)]
pub struct RedisStateManager {
    connection_string: String,
}

/// `PostgreSQL` state manager  
#[derive(Debug)]
pub struct PostgresStateManager {
    connection_string: String,
}

/// In-memory state backend
#[derive(Debug)]
struct InMemoryBackend {
    data: Arc<RwLock<HashMap<String, serde_json::Value>>>,
}

impl StateManager {
    pub async fn new(config: StateConfig) -> Result<Self> {
        let backend: Arc<dyn StateBackend> = match &config.backend {
            StateBackendType::InMemory => Arc::new(InMemoryBackend {
                data: Arc::new(RwLock::new(HashMap::new())),
            }),
            StateBackendType::Redis { connection_string } => Arc::new(RedisStateManager {
                connection_string: connection_string.clone(),
            }),
            StateBackendType::PostgreSQL { connection_string } => Arc::new(PostgresStateManager {
                connection_string: connection_string.clone(),
            }),
        };

        Ok(Self {
            config,
            backend,
            state_cache: Arc::new(RwLock::new(HashMap::new())),
        })
    }

    pub async fn get_state(&self, key: &str) -> Result<Option<serde_json::Value>> {
        self.backend.get(key).await
    }

    pub async fn set_state(&self, key: &str, value: serde_json::Value) -> Result<()> {
        self.backend.set(key, value).await
    }

    pub async fn update_task_state(&self, task_id: &str, state_value: &str) -> Result<()> {
        let key = format!("task_state_{task_id}");
        let value = serde_json::json!(state_value);
        self.set_state(&key, value).await
    }

    pub async fn create_snapshot(&self) -> Result<StateSnapshot> {
        let keys = self.backend.list_keys("").await?;
        let mut entries = HashMap::new();

        for key in keys {
            if let Some(value) = self.backend.get(&key).await? {
                entries.insert(key, value);
            }
        }

        Ok(StateSnapshot {
            snapshot_id: uuid::Uuid::new_v4().to_string(),
            timestamp: Utc::now(),
            entries,
        })
    }
}

impl StateBackend for InMemoryBackend {
    fn get(
        &self,
        key: &str,
    ) -> std::pin::Pin<
        Box<dyn std::future::Future<Output = Result<Option<serde_json::Value>>> + Send + '_>,
    > {
        let key = key.to_string();
        let data = self.data.clone();
        Box::pin(async move {
            let data = data.read().await;
            Ok(data.get(&key).cloned())
        })
    }

    fn set(
        &self,
        key: &str,
        value: serde_json::Value,
    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send + '_>> {
        let key = key.to_string();
        let data = self.data.clone();
        Box::pin(async move {
            let mut data = data.write().await;
            data.insert(key, value);
            Ok(())
        })
    }

    fn delete(
        &self,
        key: &str,
    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send + '_>> {
        let key = key.to_string();
        let data = self.data.clone();
        Box::pin(async move {
            let mut data = data.write().await;
            data.remove(&key);
            Ok(())
        })
    }

    fn list_keys(
        &self,
        prefix: &str,
    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<Vec<String>>> + Send + '_>> {
        let prefix = prefix.to_string();
        let data = self.data.clone();
        Box::pin(async move {
            let data = data.read().await;
            Ok(data
                .keys()
                .filter(|key| key.starts_with(&prefix))
                .cloned()
                .collect())
        })
    }
}

impl StateBackend for RedisStateManager {
    fn get(
        &self,
        _key: &str,
    ) -> std::pin::Pin<
        Box<dyn std::future::Future<Output = Result<Option<serde_json::Value>>> + Send + '_>,
    > {
        Box::pin(async move { Err(EtlError::Other("Redis backend not implemented".to_string())) })
    }

    fn set(
        &self,
        _key: &str,
        _value: serde_json::Value,
    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send + '_>> {
        Box::pin(async move { Err(EtlError::Other("Redis backend not implemented".to_string())) })
    }

    fn delete(
        &self,
        _key: &str,
    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send + '_>> {
        Box::pin(async move { Err(EtlError::Other("Redis backend not implemented".to_string())) })
    }

    fn list_keys(
        &self,
        _prefix: &str,
    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<Vec<String>>> + Send + '_>> {
        Box::pin(async move { Err(EtlError::Other("Redis backend not implemented".to_string())) })
    }
}

impl StateBackend for PostgresStateManager {
    fn get(
        &self,
        _key: &str,
    ) -> std::pin::Pin<
        Box<dyn std::future::Future<Output = Result<Option<serde_json::Value>>> + Send + '_>,
    > {
        Box::pin(async move {
            Err(EtlError::Other(
                "PostgreSQL backend not implemented".to_string(),
            ))
        })
    }

    fn set(
        &self,
        _key: &str,
        _value: serde_json::Value,
    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send + '_>> {
        Box::pin(async move {
            Err(EtlError::Other(
                "PostgreSQL backend not implemented".to_string(),
            ))
        })
    }

    fn delete(
        &self,
        _key: &str,
    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<()>> + Send + '_>> {
        Box::pin(async move {
            Err(EtlError::Other(
                "PostgreSQL backend not implemented".to_string(),
            ))
        })
    }

    fn list_keys(
        &self,
        _prefix: &str,
    ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<Vec<String>>> + Send + '_>> {
        Box::pin(async move {
            Err(EtlError::Other(
                "PostgreSQL backend not implemented".to_string(),
            ))
        })
    }
}