//! State management system for workflow orchestrator
//!
//! This module provides persistent state tracking, checkpointing, and recovery
//! functionality that integrates with the existing storage abstraction layer.
//! It supports both SurrealDB and filesystem backends for state persistence.

use std::{
    collections::HashMap,
    sync::Arc,
    time::{Duration, Instant},
};

use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use tokio::sync::{Mutex, RwLock};
use tracing::{debug, error, info, warn};
use uuid::Uuid;

use crate::{
    storage::{Metric, TaskStorage},
    workflow::{WorkflowExecutionState, WorkflowStatus},
};

/// Configuration for state management
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StateConfig {
    /// Checkpoint interval in seconds
    pub checkpoint_interval_seconds: u64,
    /// Maximum number of checkpoints to retain
    pub max_checkpoints: usize,
    /// Enable automatic state recovery on startup
    pub auto_recovery: bool,
    /// Cleanup interval for old state data
    pub cleanup_interval_hours: u64,
    /// State data retention period in days
    pub retention_days: u64,
}

impl Default for StateConfig {
    fn default() -> Self {
        Self {
            checkpoint_interval_seconds: 300, // 5 minutes
            max_checkpoints: 50,
            auto_recovery: true,
            cleanup_interval_hours: 24,
            retention_days: 30,
        }
    }
}

/// Workflow state snapshot for checkpointing
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StateSnapshot {
    pub snapshot_id: String,
    pub workflow_id: String,
    pub timestamp: chrono::DateTime<chrono::Utc>,
    pub workflow_state: WorkflowExecutionState,
    pub system_metrics: SystemMetrics,
    pub metadata: HashMap<String, String>,
}

/// System performance and resource metrics
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SystemMetrics {
    pub memory_usage_mb: u64,
    pub cpu_usage_percent: f64,
    pub active_tasks: usize,
    pub completed_tasks: usize,
    pub failed_tasks: usize,
    pub storage_size_mb: u64,
    pub uptime_seconds: u64,
    pub throughput_tasks_per_minute: f64,
}

impl Default for SystemMetrics {
    fn default() -> Self {
        Self {
            memory_usage_mb: 0,
            cpu_usage_percent: 0.0,
            active_tasks: 0,
            completed_tasks: 0,
            failed_tasks: 0,
            storage_size_mb: 0,
            uptime_seconds: 0,
            throughput_tasks_per_minute: 0.0,
        }
    }
}

impl SystemMetrics {
    pub fn new() -> Self {
        Self::default()
    }
}

/// Progress tracking for real-time updates
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProgressInfo {
    pub workflow_id: String,
    pub total_tasks: usize,
    pub completed_tasks: usize,
    pub failed_tasks: usize,
    pub running_tasks: usize,
    pub pending_tasks: usize,
    pub progress_percentage: f64,
    pub estimated_completion: Option<chrono::DateTime<chrono::Utc>>,
    pub current_phase: String,
}

impl ProgressInfo {
    pub fn new(workflow_id: String, total_tasks: usize) -> Self {
        Self {
            workflow_id,
            total_tasks,
            completed_tasks: 0,
            failed_tasks: 0,
            running_tasks: 0,
            pending_tasks: total_tasks,
            progress_percentage: 0.0,
            estimated_completion: None,
            current_phase: "Initializing".to_string(),
        }
    }

    pub fn calculate_progress(&mut self) {
        if self.total_tasks > 0 {
            self.progress_percentage =
                (self.completed_tasks as f64 / self.total_tasks as f64) * 100.0;
        }
    }

    pub fn estimate_completion(&mut self, average_task_duration_ms: u64) {
        if self.running_tasks > 0 && average_task_duration_ms > 0 {
            let remaining_tasks = self.pending_tasks;
            let estimated_ms = (remaining_tasks as u64) * average_task_duration_ms;
            self.estimated_completion =
                Some(chrono::Utc::now() + chrono::Duration::milliseconds(estimated_ms as i64));
        }
    }
}

/// Health check result for system monitoring
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HealthStatus {
    pub component: String,
    pub status: ComponentHealth,
    pub last_check: chrono::DateTime<chrono::Utc>,
    pub details: HashMap<String, String>,
    pub error: Option<String>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum ComponentHealth {
    Healthy,
    Degraded,
    Unhealthy,
    Unknown,
}

impl std::fmt::Display for ComponentHealth {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ComponentHealth::Healthy => write!(f, "healthy"),
            ComponentHealth::Degraded => write!(f, "degraded"),
            ComponentHealth::Unhealthy => write!(f, "unhealthy"),
            ComponentHealth::Unknown => write!(f, "unknown"),
        }
    }
}

/// Main workflow state manager
pub struct WorkflowStateManager {
    storage: Arc<dyn TaskStorage>,
    config: StateConfig,
    snapshots: Arc<RwLock<HashMap<String, Vec<StateSnapshot>>>>,
    progress_tracker: Arc<RwLock<HashMap<String, ProgressInfo>>>,
    health_status: Arc<RwLock<HashMap<String, HealthStatus>>>,
    system_metrics: Arc<RwLock<SystemMetrics>>,
    start_time: Instant,
    last_cleanup: Arc<Mutex<Instant>>,
}

impl WorkflowStateManager {
    /// Create new state manager
    pub fn new(storage: Arc<dyn TaskStorage>, config: StateConfig) -> Self {
        Self {
            storage,
            config,
            snapshots: Arc::new(RwLock::new(HashMap::new())),
            progress_tracker: Arc::new(RwLock::new(HashMap::new())),
            health_status: Arc::new(RwLock::new(HashMap::new())),
            system_metrics: Arc::new(RwLock::new(SystemMetrics::new())),
            start_time: Instant::now(),
            last_cleanup: Arc::new(Mutex::new(Instant::now())),
        }
    }

    /// Initialize state manager
    pub async fn init(&self) -> Result<()> {
        info!("Initializing workflow state manager");

        // Initialize storage backend
        self.storage
            .init()
            .await
            .context("Failed to initialize storage backend")?;

        // Perform health check
        self.update_health_status("storage", ComponentHealth::Healthy, HashMap::new(), None)
            .await;

        // Start background tasks
        self.start_background_tasks().await?;

        // Auto-recovery if enabled
        if self.config.auto_recovery {
            if let Err(e) = self.recover_workflows().await {
                warn!("Failed to recover workflows: {}", e);
            }
        }

        info!("State manager initialized successfully");
        Ok(())
    }

    /// Create workflow state checkpoint
    pub async fn create_checkpoint(
        &self,
        workflow_state: &WorkflowExecutionState,
    ) -> Result<StateSnapshot> {
        let snapshot_id = Uuid::new_v4().to_string();

        // Collect current system metrics
        let system_metrics = {
            let metrics = self.system_metrics.read().await;
            metrics.clone()
        };

        let snapshot = StateSnapshot {
            snapshot_id: snapshot_id.clone(),
            workflow_id: workflow_state.workflow_id.clone(),
            timestamp: chrono::Utc::now(),
            workflow_state: workflow_state.clone(),
            system_metrics,
            metadata: HashMap::new(),
        };

        // Store snapshot
        {
            let mut snapshots = self.snapshots.write().await;
            let workflow_snapshots = snapshots
                .entry(workflow_state.workflow_id.clone())
                .or_default();
            workflow_snapshots.push(snapshot.clone());

            // Keep only max_checkpoints
            if workflow_snapshots.len() > self.config.max_checkpoints {
                workflow_snapshots.drain(0..workflow_snapshots.len() - self.config.max_checkpoints);
            }
        }

        // Persist to storage
        let metric = Metric::counter("checkpoint_created", 1.0)
            .with_tag("workflow_id", &workflow_state.workflow_id);
        let _ = self.storage.record_metric(metric).await;

        debug!(
            "Created checkpoint {} for workflow {}",
            snapshot_id, workflow_state.workflow_id
        );
        Ok(snapshot)
    }

    /// Restore workflow from latest checkpoint
    pub async fn restore_workflow(
        &self,
        workflow_id: &str,
    ) -> Result<Option<WorkflowExecutionState>> {
        let snapshots = self.snapshots.read().await;

        if let Some(workflow_snapshots) = snapshots.get(workflow_id) {
            if let Some(latest_snapshot) = workflow_snapshots.last() {
                info!(
                    "Restoring workflow {} from checkpoint {}",
                    workflow_id, latest_snapshot.snapshot_id
                );
                return Ok(Some(latest_snapshot.workflow_state.clone()));
            }
        }

        Ok(None)
    }

    /// Update workflow progress
    pub async fn update_progress(&self, workflow_id: &str, progress: ProgressInfo) {
        let mut tracker = self.progress_tracker.write().await;
        tracker.insert(workflow_id.to_string(), progress);
    }

    /// Get workflow progress
    pub async fn get_progress(&self, workflow_id: &str) -> Option<ProgressInfo> {
        let tracker = self.progress_tracker.read().await;
        tracker.get(workflow_id).cloned()
    }

    /// List all workflow progress
    pub async fn list_progress(&self) -> HashMap<String, ProgressInfo> {
        let tracker = self.progress_tracker.read().await;
        tracker.clone()
    }

    /// Update system health status
    pub async fn update_health_status(
        &self,
        component: &str,
        status: ComponentHealth,
        details: HashMap<String, String>,
        error: Option<String>,
    ) {
        let health_status = HealthStatus {
            component: component.to_string(),
            status: status.clone(),
            last_check: chrono::Utc::now(),
            details,
            error,
        };

        let mut health = self.health_status.write().await;
        health.insert(component.to_string(), health_status);

        // Record metric
        let metric_value = match status {
            ComponentHealth::Healthy => 1.0,
            ComponentHealth::Degraded => 0.5,
            ComponentHealth::Unhealthy => 0.0,
            ComponentHealth::Unknown => -1.0,
        };

        let metric =
            Metric::gauge("component_health", metric_value).with_tag("component", component);
        let _ = self.storage.record_metric(metric).await;
    }

    /// Get system health status
    pub async fn get_health_status(&self) -> HashMap<String, HealthStatus> {
        let health = self.health_status.read().await;
        health.clone()
    }

    /// Update system metrics
    pub async fn update_system_metrics(&self, metrics: SystemMetrics) {
        {
            let mut system_metrics = self.system_metrics.write().await;
            *system_metrics = metrics.clone();
        }

        // Record individual metrics
        let metrics = vec![
            Metric::gauge("memory_usage_mb", metrics.memory_usage_mb as f64),
            Metric::gauge("cpu_usage_percent", metrics.cpu_usage_percent),
            Metric::gauge("active_tasks", metrics.active_tasks as f64),
            Metric::counter("completed_tasks", metrics.completed_tasks as f64),
            Metric::counter("failed_tasks", metrics.failed_tasks as f64),
            Metric::gauge("storage_size_mb", metrics.storage_size_mb as f64),
            Metric::gauge("uptime_seconds", metrics.uptime_seconds as f64),
            Metric::gauge(
                "throughput_tasks_per_minute",
                metrics.throughput_tasks_per_minute,
            ),
        ];

        for metric in metrics {
            let _ = self.storage.record_metric(metric).await;
        }
    }

    /// Get current system metrics
    pub async fn get_system_metrics(&self) -> SystemMetrics {
        let mut metrics = self.system_metrics.read().await.clone();
        metrics.uptime_seconds = self.start_time.elapsed().as_secs();
        metrics
    }

    /// Clean up old state data
    pub async fn cleanup_old_data(&self) -> Result<usize> {
        let cleanup_time =
            chrono::Utc::now() - chrono::Duration::days(self.config.retention_days as i64);
        let mut cleaned_count = 0;

        // Clean up old snapshots
        {
            let mut snapshots = self.snapshots.write().await;
            for (_workflow_id, workflow_snapshots) in snapshots.iter_mut() {
                let initial_count = workflow_snapshots.len();
                workflow_snapshots.retain(|snapshot| snapshot.timestamp > cleanup_time);
                cleaned_count += initial_count - workflow_snapshots.len();
            }

            // Remove empty entries
            snapshots.retain(|_, snapshots| !snapshots.is_empty());
        }

        // Clean up old progress data
        {
            let mut tracker = self.progress_tracker.write().await;
            // For now, just keep active workflows - in production might want more
            // sophisticated cleanup
            tracker.retain(|_workflow_id, progress| {
                progress.completed_tasks + progress.failed_tasks < progress.total_tasks
            });
        }

        if cleaned_count > 0 {
            info!("Cleaned up {} old state entries", cleaned_count);
            let metric = Metric::counter("state_cleanup_count", cleaned_count as f64);
            let _ = self.storage.record_metric(metric).await;
        }

        Ok(cleaned_count)
    }

    /// Get workflow state snapshots
    pub async fn get_workflow_snapshots(&self, workflow_id: &str) -> Vec<StateSnapshot> {
        let snapshots = self.snapshots.read().await;
        snapshots.get(workflow_id).cloned().unwrap_or_default()
    }

    /// Get system statistics
    pub async fn get_statistics(&self) -> Result<StateManagerStatistics> {
        let snapshots = self.snapshots.read().await;
        let progress = self.progress_tracker.read().await;
        let health = self.health_status.read().await;
        let system_metrics = self.get_system_metrics().await;

        let total_snapshots: usize = snapshots.values().map(|s| s.len()).sum();
        let active_workflows = progress.len();
        let healthy_components = health
            .values()
            .filter(|h| h.status == ComponentHealth::Healthy)
            .count();

        Ok(StateManagerStatistics {
            total_snapshots,
            active_workflows,
            total_components: health.len(),
            healthy_components,
            system_metrics,
            uptime_seconds: self.start_time.elapsed().as_secs(),
        })
    }

    /// Start background maintenance tasks
    async fn start_background_tasks(&self) -> Result<()> {
        let state_manager = self.clone_arc();

        // Cleanup task
        tokio::spawn(async move {
            let mut interval = tokio::time::interval(Duration::from_secs(
                state_manager.config.cleanup_interval_hours * 3600,
            ));

            loop {
                interval.tick().await;

                if let Err(e) = state_manager.cleanup_old_data().await {
                    error!("Failed to cleanup old state data: {}", e);
                }

                // Update last cleanup time
                {
                    let mut last_cleanup = state_manager.last_cleanup.lock().await;
                    *last_cleanup = Instant::now();
                }
            }
        });

        Ok(())
    }

    /// Recover workflows from checkpoints
    async fn recover_workflows(&self) -> Result<Vec<String>> {
        info!("Attempting workflow recovery from checkpoints");
        let mut recovered_workflows = Vec::new();

        let snapshots = self.snapshots.read().await;
        for (workflow_id, workflow_snapshots) in snapshots.iter() {
            if let Some(latest_snapshot) = workflow_snapshots.last() {
                if latest_snapshot.workflow_state.status == WorkflowStatus::Running {
                    // In a full implementation, this would restart the workflow
                    info!("Found recoverable workflow: {}", workflow_id);
                    recovered_workflows.push(workflow_id.clone());
                }
            }
        }

        if !recovered_workflows.is_empty() {
            info!(
                "Recovered {} workflows from checkpoints",
                recovered_workflows.len()
            );
        }

        Ok(recovered_workflows)
    }

    /// Helper method to clone Arc-wrapped self for background tasks
    fn clone_arc(&self) -> Arc<Self> {
        Arc::new(Self {
            storage: self.storage.clone(),
            config: self.config.clone(),
            snapshots: self.snapshots.clone(),
            progress_tracker: self.progress_tracker.clone(),
            health_status: self.health_status.clone(),
            system_metrics: self.system_metrics.clone(),
            start_time: self.start_time,
            last_cleanup: self.last_cleanup.clone(),
        })
    }
}

/// Statistics for state manager
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StateManagerStatistics {
    pub total_snapshots: usize,
    pub active_workflows: usize,
    pub total_components: usize,
    pub healthy_components: usize,
    pub system_metrics: SystemMetrics,
    pub uptime_seconds: u64,
}

/// Progress tracker for real-time workflow monitoring
pub struct ProgressTracker {
    state_manager: Arc<WorkflowStateManager>,
    update_interval: Duration,
}

impl ProgressTracker {
    /// Create new progress tracker
    pub fn new(state_manager: Arc<WorkflowStateManager>) -> Self {
        Self {
            state_manager,
            update_interval: Duration::from_secs(5),
        }
    }

    /// Start progress tracking for a workflow
    pub async fn start_tracking(&self, workflow_id: String, total_tasks: usize) -> Result<()> {
        let progress = ProgressInfo::new(workflow_id.clone(), total_tasks);
        self.state_manager
            .update_progress(&workflow_id, progress)
            .await;

        info!("Started progress tracking for workflow: {}", workflow_id);
        Ok(())
    }

    /// Update task completion progress
    pub async fn update_task_completion(
        &self,
        workflow_id: &str,
        completed: usize,
        failed: usize,
        running: usize,
    ) -> Result<()> {
        if let Some(mut progress) = self.state_manager.get_progress(workflow_id).await {
            progress.completed_tasks = completed;
            progress.failed_tasks = failed;
            progress.running_tasks = running;
            progress.pending_tasks = progress.total_tasks - completed - failed - running;
            progress.calculate_progress();

            // Update phase based on progress
            progress.current_phase = if progress.progress_percentage == 0.0 {
                "Starting".to_string()
            } else if progress.progress_percentage < 25.0 {
                "Early Phase".to_string()
            } else if progress.progress_percentage < 75.0 {
                "Mid Phase".to_string()
            } else if progress.progress_percentage < 100.0 {
                "Final Phase".to_string()
            } else {
                "Completed".to_string()
            };

            self.state_manager
                .update_progress(workflow_id, progress)
                .await;
        }

        Ok(())
    }

    /// Get real-time progress for a workflow
    pub async fn get_real_time_progress(&self, workflow_id: &str) -> Option<ProgressInfo> {
        self.state_manager.get_progress(workflow_id).await
    }
}
