//! Rollback and recovery system with storage abstraction support
//!
//! This module provides checkpoint/rollback functionality that integrates with
//! the existing storage backends and supports provider-specific cleanup
//! operations.

use std::{
    collections::{HashMap, HashSet, VecDeque},
    fmt,
    sync::Arc,
    time::Instant,
};

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

use crate::{
    storage::{AuditEntry, Metric, TaskStorage},
    workflow::WorkflowExecutionState,
    TaskStatus, WorkflowTask,
};

/// Configuration for rollback system
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RollbackConfig {
    /// Checkpoint interval in seconds
    pub checkpoint_interval_seconds: u64,
    /// Maximum number of checkpoints to retain
    pub max_checkpoints: usize,
    /// Enable automatic checkpointing
    pub auto_checkpoint_enabled: bool,
    /// Rollback timeout in seconds
    pub rollback_timeout_seconds: u64,
    /// Resource cleanup timeout in seconds
    pub cleanup_timeout_seconds: u64,
    /// Enable partial rollbacks
    pub partial_rollback_enabled: bool,
    /// Rollback strategy
    pub strategy: RollbackStrategy,
}

impl Default for RollbackConfig {
    fn default() -> Self {
        Self {
            checkpoint_interval_seconds: 300, // 5 minutes
            max_checkpoints: 20,
            auto_checkpoint_enabled: true,
            rollback_timeout_seconds: 1800, // 30 minutes
            cleanup_timeout_seconds: 600,   // 10 minutes
            partial_rollback_enabled: true,
            strategy: RollbackStrategy::ConfigDriven,
        }
    }
}

/// Rollback strategy configuration
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum RollbackStrategy {
    /// Configuration-driven strategy - use config to determine rollback
    /// approach
    ConfigDriven,
    /// Conservative strategy - preserve resources unless explicitly marked for
    /// deletion
    Conservative,
    /// Aggressive strategy - rollback all changes
    Aggressive,
    /// Custom strategy with specific operations
    Custom { operations: Vec<String> },
}

/// Checkpoint containing system state at a point in time
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Checkpoint {
    pub id: String,
    pub name: String,
    pub description: Option<String>,
    pub created_at: DateTime<Utc>,
    pub workflow_state: Option<WorkflowExecutionState>,
    pub resources: Vec<ResourceSnapshot>,
    pub metadata: HashMap<String, String>,
    pub provider_states: HashMap<String, ProviderState>,
    pub task_states: Vec<WorkflowTask>,
    pub dependencies: Vec<String>,
}

impl Checkpoint {
    /// Create a new checkpoint with minimal information
    pub fn new(name: String) -> Self {
        Self {
            id: Uuid::new_v4().to_string(),
            name,
            description: None,
            created_at: Utc::now(),
            workflow_state: None,
            resources: Vec::new(),
            metadata: HashMap::new(),
            provider_states: HashMap::new(),
            task_states: Vec::new(),
            dependencies: Vec::new(),
        }
    }

    /// Add metadata to the checkpoint
    pub fn with_metadata(mut self, key: String, value: String) -> Self {
        self.metadata.insert(key, value);
        self
    }

    /// Add provider state to the checkpoint
    pub fn with_provider_state(mut self, provider: String, state: ProviderState) -> Self {
        self.provider_states.insert(provider, state);
        self
    }

    /// Check if checkpoint is expired based on retention policy
    pub fn is_expired(&self, retention_days: i64) -> bool {
        let expiry_date = self.created_at + Duration::days(retention_days);
        Utc::now() > expiry_date
    }
}

/// Snapshot of a resource at checkpoint time
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResourceSnapshot {
    pub resource_id: String,
    pub resource_type: String,
    pub provider: String,
    pub state: ResourceState,
    pub created_at: DateTime<Utc>,
    pub dependencies: Vec<String>,
    pub config: HashMap<String, String>,
    pub tags: HashMap<String, String>,
}

/// Current state of a resource
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum ResourceState {
    /// Resource exists and is operational
    Active,
    /// Resource is being created
    Creating,
    /// Resource is being modified
    Updating,
    /// Resource is being deleted
    Deleting,
    /// Resource has been deleted
    Deleted,
    /// Resource is in error state
    Error { reason: String },
    /// Resource state is unknown
    Unknown,
}

impl fmt::Display for ResourceState {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            ResourceState::Active => write!(f, "active"),
            ResourceState::Creating => write!(f, "creating"),
            ResourceState::Updating => write!(f, "updating"),
            ResourceState::Deleting => write!(f, "deleting"),
            ResourceState::Deleted => write!(f, "deleted"),
            ResourceState::Error { reason } => write!(f, "error: {}", reason),
            ResourceState::Unknown => write!(f, "unknown"),
        }
    }
}

/// Provider-specific state information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ProviderState {
    pub provider_type: String,
    pub region: Option<String>,
    pub config: HashMap<String, String>,
    pub resources: Vec<String>,
    pub active_operations: Vec<String>,
    pub auth_state: Option<String>,
}

/// Rollback operation result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RollbackResult {
    pub success: bool,
    pub checkpoint_id: String,
    pub operations_executed: usize,
    pub operations_failed: usize,
    pub rolled_back_resources: Vec<String>,
    pub failed_resources: Vec<String>,
    pub duration_seconds: f64,
    pub errors: Vec<String>,
    pub warnings: Vec<String>,
}

impl RollbackResult {
    /// Create a successful rollback result
    pub fn success(checkpoint_id: String, duration: f64) -> Self {
        Self {
            success: true,
            checkpoint_id,
            operations_executed: 0,
            operations_failed: 0,
            rolled_back_resources: Vec::new(),
            failed_resources: Vec::new(),
            duration_seconds: duration,
            errors: Vec::new(),
            warnings: Vec::new(),
        }
    }

    /// Create a failed rollback result
    pub fn failure(checkpoint_id: String, error: String, duration: f64) -> Self {
        Self {
            success: false,
            checkpoint_id,
            operations_executed: 0,
            operations_failed: 1,
            rolled_back_resources: Vec::new(),
            failed_resources: Vec::new(),
            duration_seconds: duration,
            errors: vec![error],
            warnings: Vec::new(),
        }
    }

    /// Add an executed operation
    pub fn add_operation(&mut self, success: bool) {
        if success {
            self.operations_executed += 1;
        } else {
            self.operations_failed += 1;
        }
    }

    /// Add a rolled back resource
    pub fn add_rolled_back_resource(&mut self, resource_id: String) {
        self.rolled_back_resources.push(resource_id);
    }

    /// Add a failed resource
    pub fn add_failed_resource(&mut self, resource_id: String) {
        self.failed_resources.push(resource_id);
    }

    /// Add an error
    pub fn add_error(&mut self, error: String) {
        self.errors.push(error);
    }

    /// Add a warning
    pub fn add_warning(&mut self, warning: String) {
        self.warnings.push(warning);
    }
}

/// Trait for provider-specific rollback handlers
#[async_trait]
pub trait ProviderRollbackHandler: Send + Sync {
    /// Provider type this handler supports
    fn provider_type(&self) -> &str;

    /// Rollback a specific resource
    async fn rollback_resource(&self, resource: &ResourceSnapshot) -> Result<()>;

    /// Validate that a resource can be rolled back
    async fn can_rollback_resource(&self, resource: &ResourceSnapshot) -> Result<bool>;

    /// Get current state of a resource
    async fn get_resource_state(&self, resource_id: &str) -> Result<ResourceState>;

    /// Cleanup resources created after a checkpoint
    async fn cleanup_resources_after(&self, checkpoint: &Checkpoint) -> Result<Vec<String>>;

    /// Validate rollback before execution
    async fn validate_rollback(&self, checkpoint: &Checkpoint) -> Result<Vec<String>>;
}

/// Manages resource tracking and state snapshots
pub struct ResourceTracker {
    storage: Arc<dyn TaskStorage>,
    tracked_resources: Arc<RwLock<HashMap<String, ResourceSnapshot>>>,
    provider_handlers: Arc<RwLock<HashMap<String, Arc<dyn ProviderRollbackHandler>>>>,
}

impl ResourceTracker {
    /// Create a new resource tracker
    pub fn new(storage: Arc<dyn TaskStorage>) -> Self {
        Self {
            storage,
            tracked_resources: Arc::new(RwLock::new(HashMap::new())),
            provider_handlers: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Register a provider-specific rollback handler
    pub async fn register_provider_handler(&self, handler: Arc<dyn ProviderRollbackHandler>) {
        let provider_type = handler.provider_type().to_string();
        let mut handlers = self.provider_handlers.write().await;
        handlers.insert(provider_type.clone(), handler);
        info!(
            "Registered rollback handler for provider: {}",
            provider_type
        );
    }

    /// Track a new resource
    pub async fn track_resource(&self, resource: ResourceSnapshot) -> Result<()> {
        let resource_id = resource.resource_id.clone();

        // Store in memory cache
        {
            let mut resources = self.tracked_resources.write().await;
            resources.insert(resource_id.clone(), resource.clone());
        }

        // Record audit entry
        let audit_entry = AuditEntry {
            id: Uuid::new_v4().to_string(),
            task_id: resource_id.clone(),
            operation: "resource_tracked".to_string(),
            old_status: None,
            new_status: None,
            user_id: None,
            timestamp: Utc::now(),
            metadata: {
                let mut meta = HashMap::new();
                meta.insert("resource_type".to_string(), resource.resource_type);
                meta.insert("provider".to_string(), resource.provider);
                meta.insert("state".to_string(), resource.state.to_string());
                meta
            },
        };

        self.storage
            .record_audit_entry(audit_entry)
            .await
            .context("Failed to record resource tracking audit entry")?;

        debug!("Started tracking resource: {}", resource_id);
        Ok(())
    }

    /// Update resource state
    pub async fn update_resource_state(
        &self,
        resource_id: &str,
        new_state: ResourceState,
    ) -> Result<()> {
        let old_state = {
            let mut resources = self.tracked_resources.write().await;
            if let Some(resource) = resources.get_mut(resource_id) {
                let old_state = resource.state.clone();
                resource.state = new_state.clone();
                old_state
            } else {
                return Err(anyhow::anyhow!("Resource not found: {}", resource_id));
            }
        };

        // Record state change
        let audit_entry = AuditEntry {
            id: Uuid::new_v4().to_string(),
            task_id: resource_id.to_string(),
            operation: "resource_state_changed".to_string(),
            old_status: None,
            new_status: None,
            user_id: None,
            timestamp: Utc::now(),
            metadata: {
                let mut meta = HashMap::new();
                meta.insert("old_state".to_string(), old_state.to_string());
                meta.insert("new_state".to_string(), new_state.to_string());
                meta
            },
        };

        self.storage
            .record_audit_entry(audit_entry)
            .await
            .context("Failed to record resource state change")?;

        debug!(
            "Updated resource {} state: {} -> {}",
            resource_id, old_state, new_state
        );
        Ok(())
    }

    /// Get all tracked resources
    pub async fn get_tracked_resources(&self) -> HashMap<String, ResourceSnapshot> {
        self.tracked_resources.read().await.clone()
    }

    /// Get resources created after a specific checkpoint
    pub async fn get_resources_after_checkpoint(
        &self,
        checkpoint: &Checkpoint,
    ) -> Vec<ResourceSnapshot> {
        let resources = self.tracked_resources.read().await;
        resources
            .values()
            .filter(|resource| resource.created_at > checkpoint.created_at)
            .cloned()
            .collect()
    }

    /// Get provider handler for a resource
    pub async fn get_provider_handler(
        &self,
        provider: &str,
    ) -> Option<Arc<dyn ProviderRollbackHandler>> {
        self.provider_handlers.read().await.get(provider).cloned()
    }
}

/// Manages checkpoints and system state snapshots
pub struct CheckpointManager {
    storage: Arc<dyn TaskStorage>,
    config: RollbackConfig,
    checkpoints: Arc<RwLock<VecDeque<Checkpoint>>>,
    last_checkpoint: Arc<RwLock<Option<DateTime<Utc>>>>,
}

impl CheckpointManager {
    /// Create a new checkpoint manager
    pub fn new(storage: Arc<dyn TaskStorage>, config: RollbackConfig) -> Self {
        Self {
            storage,
            config,
            checkpoints: Arc::new(RwLock::new(VecDeque::new())),
            last_checkpoint: Arc::new(RwLock::new(None)),
        }
    }

    /// Create a new checkpoint
    pub async fn create_checkpoint(
        &self,
        name: String,
        description: Option<String>,
    ) -> Result<String> {
        let start_time = Instant::now();

        // Get current system state
        let tasks = self
            .storage
            .list_tasks(None)
            .await
            .context("Failed to get current tasks for checkpoint")?;

        let mut checkpoint = Checkpoint::new(name.clone());
        checkpoint.description = description.clone();
        checkpoint.task_states = tasks;

        // Add system metrics
        checkpoint.metadata.insert(
            "total_tasks".to_string(),
            self.storage.total_tasks().await.unwrap_or(0).to_string(),
        );
        checkpoint.metadata.insert(
            "queue_size".to_string(),
            self.storage.queue_size().await.unwrap_or(0).to_string(),
        );

        // Store checkpoint
        {
            let mut checkpoints = self.checkpoints.write().await;
            checkpoints.push_back(checkpoint.clone());

            // Maintain checkpoint limit
            while checkpoints.len() > self.config.max_checkpoints {
                if let Some(old_checkpoint) = checkpoints.pop_front() {
                    info!(
                        "Removed old checkpoint: {} ({})",
                        old_checkpoint.id, old_checkpoint.name
                    );
                }
            }
        }

        // Update last checkpoint time
        {
            let mut last = self.last_checkpoint.write().await;
            *last = Some(checkpoint.created_at);
        }

        // Record metrics
        let duration = start_time.elapsed().as_secs_f64();
        let metric = Metric::counter("checkpoint_created", 1.0)
            .with_tag("name", &name)
            .with_tag("duration_seconds", &duration.to_string());

        self.storage
            .record_metric(metric)
            .await
            .context("Failed to record checkpoint creation metric")?;

        // Record audit entry
        let audit_entry = AuditEntry {
            id: Uuid::new_v4().to_string(),
            task_id: checkpoint.id.clone(),
            operation: "checkpoint_created".to_string(),
            old_status: None,
            new_status: None,
            user_id: None,
            timestamp: checkpoint.created_at,
            metadata: {
                let mut meta = HashMap::new();
                meta.insert("name".to_string(), name.clone());
                meta.insert(
                    "task_count".to_string(),
                    checkpoint.task_states.len().to_string(),
                );
                if let Some(desc) = &description {
                    meta.insert("description".to_string(), desc.clone());
                }
                meta
            },
        };

        self.storage
            .record_audit_entry(audit_entry)
            .await
            .context("Failed to record checkpoint audit entry")?;

        info!(
            "Created checkpoint '{}' with ID: {} ({} tasks, {:.2}s)",
            name,
            checkpoint.id,
            checkpoint.task_states.len(),
            duration
        );
        Ok(checkpoint.id)
    }

    /// Get a checkpoint by ID
    pub async fn get_checkpoint(&self, id: &str) -> Option<Checkpoint> {
        let checkpoints = self.checkpoints.read().await;
        checkpoints.iter().find(|cp| cp.id == id).cloned()
    }

    /// List all checkpoints
    pub async fn list_checkpoints(&self) -> Vec<Checkpoint> {
        let checkpoints = self.checkpoints.read().await;
        checkpoints.iter().cloned().collect()
    }

    /// Get the latest checkpoint
    pub async fn get_latest_checkpoint(&self) -> Option<Checkpoint> {
        let checkpoints = self.checkpoints.read().await;
        checkpoints.back().cloned()
    }

    /// Check if it's time for an automatic checkpoint
    pub async fn should_create_checkpoint(&self) -> bool {
        if !self.config.auto_checkpoint_enabled {
            return false;
        }

        let last_checkpoint = self.last_checkpoint.read().await;
        match *last_checkpoint {
            Some(last_time) => {
                let elapsed = Utc::now().signed_duration_since(last_time);
                elapsed.num_seconds() >= self.config.checkpoint_interval_seconds as i64
            }
            None => true, // No checkpoint exists yet
        }
    }

    /// Clean up expired checkpoints
    pub async fn cleanup_expired_checkpoints(&self, retention_days: i64) -> Result<usize> {
        let mut removed_count = 0;
        let mut checkpoints = self.checkpoints.write().await;

        // Use drain_filter when it becomes stable, for now collect indices to remove
        let mut indices_to_remove = Vec::new();
        for (i, checkpoint) in checkpoints.iter().enumerate() {
            if checkpoint.is_expired(retention_days) {
                indices_to_remove.push(i);
            }
        }

        // Remove in reverse order to maintain indices
        for &index in indices_to_remove.iter().rev() {
            if let Some(checkpoint) = checkpoints.get(index) {
                info!(
                    "Removing expired checkpoint: {} ({})",
                    checkpoint.id, checkpoint.name
                );
                removed_count += 1;
            }
            checkpoints.remove(index);
        }

        if removed_count > 0 {
            let metric = Metric::counter("checkpoints_cleaned", removed_count as f64);
            self.storage
                .record_metric(metric)
                .await
                .context("Failed to record checkpoint cleanup metric")?;
        }

        Ok(removed_count)
    }
}

/// Executes rollback operations
pub struct RollbackExecutor {
    storage: Arc<dyn TaskStorage>,
    resource_tracker: Arc<ResourceTracker>,
    checkpoint_manager: Arc<CheckpointManager>,
    config: RollbackConfig,
    active_rollbacks: Arc<Mutex<HashSet<String>>>,
}

impl RollbackExecutor {
    /// Create a new rollback executor
    pub fn new(
        storage: Arc<dyn TaskStorage>,
        resource_tracker: Arc<ResourceTracker>,
        checkpoint_manager: Arc<CheckpointManager>,
        config: RollbackConfig,
    ) -> Self {
        Self {
            storage,
            resource_tracker,
            checkpoint_manager,
            config,
            active_rollbacks: Arc::new(Mutex::new(HashSet::new())),
        }
    }

    /// Execute rollback to a specific checkpoint
    pub async fn rollback_to_checkpoint(&self, checkpoint_id: &str) -> Result<RollbackResult> {
        let start_time = Instant::now();

        // Check if rollback is already in progress
        {
            let mut active = self.active_rollbacks.lock().await;
            if active.contains(checkpoint_id) {
                return Err(anyhow::anyhow!(
                    "Rollback to checkpoint {} is already in progress",
                    checkpoint_id
                ));
            }
            active.insert(checkpoint_id.to_string());
        }

        // Track that we're starting rollback for cleanup purposes
        let checkpoint_id_for_cleanup = checkpoint_id.to_string();

        let mut result = RollbackResult::success(checkpoint_id.to_string(), 0.0);

        // Get the target checkpoint
        let checkpoint = match self.checkpoint_manager.get_checkpoint(checkpoint_id).await {
            Some(cp) => cp,
            None => {
                let duration = start_time.elapsed().as_secs_f64();
                return Ok(RollbackResult::failure(
                    checkpoint_id.to_string(),
                    format!("Checkpoint not found: {}", checkpoint_id),
                    duration,
                ));
            }
        };

        info!(
            "Starting rollback to checkpoint '{}' ({})",
            checkpoint.name, checkpoint_id
        );

        // Validate rollback
        if let Err(e) = self.validate_rollback(&checkpoint).await {
            let duration = start_time.elapsed().as_secs_f64();
            result.add_error(format!("Rollback validation failed: {}", e));
            return Ok(RollbackResult::failure(
                checkpoint_id.to_string(),
                e.to_string(),
                duration,
            ));
        }

        // Execute rollback based on strategy
        match &self.config.strategy {
            RollbackStrategy::ConfigDriven => {
                self.execute_config_driven_rollback(&checkpoint, &mut result)
                    .await?;
            }
            RollbackStrategy::Conservative => {
                self.execute_conservative_rollback(&checkpoint, &mut result)
                    .await?;
            }
            RollbackStrategy::Aggressive => {
                self.execute_aggressive_rollback(&checkpoint, &mut result)
                    .await?;
            }
            RollbackStrategy::Custom { operations } => {
                self.execute_custom_rollback(&checkpoint, operations, &mut result)
                    .await?;
            }
        }

        // Cleanup active rollback tracking
        {
            let mut active = self.active_rollbacks.lock().await;
            active.remove(&checkpoint_id_for_cleanup);
        }

        // Update result with final timing
        result.duration_seconds = start_time.elapsed().as_secs_f64();
        result.success = result.operations_failed == 0 && result.errors.is_empty();

        // Record rollback metrics
        let metric = Metric::counter("rollback_executed", 1.0)
            .with_tag("checkpoint_id", checkpoint_id)
            .with_tag("success", &result.success.to_string())
            .with_tag("duration_seconds", &result.duration_seconds.to_string());

        self.storage
            .record_metric(metric)
            .await
            .context("Failed to record rollback metric")?;

        // Record audit entry
        let audit_entry = AuditEntry {
            id: Uuid::new_v4().to_string(),
            task_id: checkpoint_id.to_string(),
            operation: "rollback_executed".to_string(),
            old_status: None,
            new_status: None,
            user_id: None,
            timestamp: Utc::now(),
            metadata: {
                let mut meta = HashMap::new();
                meta.insert("success".to_string(), result.success.to_string());
                meta.insert(
                    "operations_executed".to_string(),
                    result.operations_executed.to_string(),
                );
                meta.insert(
                    "operations_failed".to_string(),
                    result.operations_failed.to_string(),
                );
                meta.insert(
                    "duration_seconds".to_string(),
                    result.duration_seconds.to_string(),
                );
                if !result.errors.is_empty() {
                    meta.insert("errors".to_string(), result.errors.join("; "));
                }
                meta
            },
        };

        self.storage
            .record_audit_entry(audit_entry)
            .await
            .context("Failed to record rollback audit entry")?;

        if result.success {
            info!(
                "Rollback completed successfully in {:.2}s: {} operations executed",
                result.duration_seconds, result.operations_executed
            );
        } else {
            error!(
                "Rollback completed with errors in {:.2}s: {}/{} operations failed",
                result.duration_seconds,
                result.operations_failed,
                result.operations_executed + result.operations_failed
            );
        }

        Ok(result)
    }

    /// Execute partial rollback for specific operations
    pub async fn rollback_operations(&self, operation_ids: Vec<String>) -> Result<RollbackResult> {
        if !self.config.partial_rollback_enabled {
            return Err(anyhow::anyhow!(
                "Partial rollback is disabled in configuration"
            ));
        }

        let start_time = Instant::now();
        let mut result = RollbackResult::success("partial".to_string(), 0.0);

        info!(
            "Starting partial rollback for {} operations",
            operation_ids.len()
        );

        // Get tasks for specified operations
        let mut rollback_tasks = Vec::new();
        for id in operation_ids {
            match self.storage.get_task(&id).await {
                Ok(Some(task)) => rollback_tasks.push(task),
                Ok(None) => {
                    result.add_warning(format!("Task not found: {}", id));
                }
                Err(e) => {
                    result.add_error(format!("Failed to retrieve task {}: {}", id, e));
                }
            }
        }

        // Execute rollback for each task
        for task in rollback_tasks {
            match self.rollback_single_task(&task).await {
                Ok(_) => {
                    result.add_operation(true);
                    result.add_rolled_back_resource(task.id.clone());
                    info!("Successfully rolled back task: {}", task.id);
                }
                Err(e) => {
                    result.add_operation(false);
                    result.add_failed_resource(task.id.clone());
                    result.add_error(format!("Failed to rollback task {}: {}", task.id, e));
                    error!("Failed to rollback task {}: {}", task.id, e);
                }
            }
        }

        result.duration_seconds = start_time.elapsed().as_secs_f64();
        result.success = result.operations_failed == 0 && result.errors.is_empty();

        info!(
            "Partial rollback completed in {:.2}s: {}/{} operations successful",
            result.duration_seconds,
            result.operations_executed,
            result.operations_executed + result.operations_failed
        );

        Ok(result)
    }

    /// Validate that rollback can be performed
    async fn validate_rollback(&self, checkpoint: &Checkpoint) -> Result<()> {
        // Check if checkpoint is too old
        let age = Utc::now().signed_duration_since(checkpoint.created_at);
        if age.num_days() > 30 {
            return Err(anyhow::anyhow!(
                "Checkpoint is too old (> 30 days): {}",
                checkpoint.name
            ));
        }

        // Validate provider states
        for provider in checkpoint.provider_states.keys() {
            if let Some(handler) = self.resource_tracker.get_provider_handler(provider).await {
                let validation_errors = handler
                    .validate_rollback(checkpoint)
                    .await
                    .context("Failed to validate provider rollback")?;

                if !validation_errors.is_empty() {
                    return Err(anyhow::anyhow!(
                        "Provider validation failed for {}: {}",
                        provider,
                        validation_errors.join(", ")
                    ));
                }
            }
        }

        // Check for conflicting active operations
        let active_rollbacks = self.active_rollbacks.lock().await;
        if !active_rollbacks.is_empty() {
            warn!(
                "Other rollback operations are active: {}",
                active_rollbacks.len()
            );
        }

        Ok(())
    }

    /// Execute config-driven rollback strategy
    async fn execute_config_driven_rollback(
        &self,
        checkpoint: &Checkpoint,
        result: &mut RollbackResult,
    ) -> Result<()> {
        info!("Executing config-driven rollback strategy");

        // Rollback tasks that were created after checkpoint
        self.rollback_tasks_after_checkpoint(checkpoint, result)
            .await?;

        // Rollback resources using provider handlers
        self.rollback_resources_after_checkpoint(checkpoint, result)
            .await?;

        Ok(())
    }

    /// Execute conservative rollback strategy
    async fn execute_conservative_rollback(
        &self,
        checkpoint: &Checkpoint,
        result: &mut RollbackResult,
    ) -> Result<()> {
        info!("Executing conservative rollback strategy");

        // Only rollback failed tasks and resources explicitly marked for deletion
        let tasks = self
            .storage
            .list_tasks(Some(TaskStatus::Failed))
            .await
            .context("Failed to get failed tasks")?;

        for task in tasks {
            if task.created_at > checkpoint.created_at {
                match self.rollback_single_task(&task).await {
                    Ok(_) => {
                        result.add_operation(true);
                        result.add_rolled_back_resource(task.id.clone());
                    }
                    Err(e) => {
                        result.add_operation(false);
                        result.add_failed_resource(task.id.clone());
                        result.add_error(format!("Failed to rollback task {}: {}", task.id, e));
                    }
                }
            }
        }

        Ok(())
    }

    /// Execute aggressive rollback strategy
    async fn execute_aggressive_rollback(
        &self,
        checkpoint: &Checkpoint,
        result: &mut RollbackResult,
    ) -> Result<()> {
        info!("Executing aggressive rollback strategy");

        // Rollback all changes made after checkpoint
        self.rollback_tasks_after_checkpoint(checkpoint, result)
            .await?;
        self.rollback_resources_after_checkpoint(checkpoint, result)
            .await?;

        // Cancel any running operations
        let running_tasks = self
            .storage
            .list_tasks(Some(TaskStatus::Running))
            .await
            .context("Failed to get running tasks")?;

        for task in running_tasks {
            if task.created_at > checkpoint.created_at {
                match self
                    .storage
                    .update_task_status(&task.id, TaskStatus::Cancelled)
                    .await
                {
                    Ok(_) => {
                        result.add_operation(true);
                        info!("Cancelled running task: {}", task.id);
                    }
                    Err(e) => {
                        result.add_operation(false);
                        result.add_error(format!("Failed to cancel task {}: {}", task.id, e));
                    }
                }
            }
        }

        Ok(())
    }

    /// Execute custom rollback strategy
    async fn execute_custom_rollback(
        &self,
        checkpoint: &Checkpoint,
        operations: &[String],
        result: &mut RollbackResult,
    ) -> Result<()> {
        info!(
            "Executing custom rollback strategy with {} operations",
            operations.len()
        );

        for operation in operations {
            match operation.as_str() {
                "rollback_tasks" => {
                    self.rollback_tasks_after_checkpoint(checkpoint, result)
                        .await?
                }
                "rollback_resources" => {
                    self.rollback_resources_after_checkpoint(checkpoint, result)
                        .await?
                }
                "cancel_running" => {
                    let running_tasks = self
                        .storage
                        .list_tasks(Some(TaskStatus::Running))
                        .await
                        .context("Failed to get running tasks")?;
                    for task in running_tasks {
                        if task.created_at > checkpoint.created_at {
                            let _ = self
                                .storage
                                .update_task_status(&task.id, TaskStatus::Cancelled)
                                .await;
                        }
                    }
                }
                _ => {
                    result.add_warning(format!("Unknown rollback operation: {}", operation));
                }
            }
        }

        Ok(())
    }

    /// Rollback tasks created after checkpoint
    async fn rollback_tasks_after_checkpoint(
        &self,
        checkpoint: &Checkpoint,
        result: &mut RollbackResult,
    ) -> Result<()> {
        let tasks = self
            .storage
            .list_tasks(None)
            .await
            .context("Failed to get tasks for rollback")?;

        let tasks_to_rollback: Vec<_> = tasks
            .into_iter()
            .filter(|task| task.created_at > checkpoint.created_at)
            .collect();

        info!(
            "Rolling back {} tasks created after checkpoint",
            tasks_to_rollback.len()
        );

        for task in tasks_to_rollback {
            match self.rollback_single_task(&task).await {
                Ok(_) => {
                    result.add_operation(true);
                    result.add_rolled_back_resource(task.id.clone());
                }
                Err(e) => {
                    result.add_operation(false);
                    result.add_failed_resource(task.id.clone());
                    result.add_error(format!("Failed to rollback task {}: {}", task.id, e));
                }
            }
        }

        Ok(())
    }

    /// Rollback resources created after checkpoint
    async fn rollback_resources_after_checkpoint(
        &self,
        checkpoint: &Checkpoint,
        result: &mut RollbackResult,
    ) -> Result<()> {
        let resources_to_rollback = self
            .resource_tracker
            .get_resources_after_checkpoint(checkpoint)
            .await;

        info!(
            "Rolling back {} resources created after checkpoint",
            resources_to_rollback.len()
        );

        for resource in resources_to_rollback {
            if let Some(handler) = self
                .resource_tracker
                .get_provider_handler(&resource.provider)
                .await
            {
                match handler.rollback_resource(&resource).await {
                    Ok(_) => {
                        result.add_operation(true);
                        result.add_rolled_back_resource(resource.resource_id.clone());
                        info!(
                            "Successfully rolled back resource: {}",
                            resource.resource_id
                        );
                    }
                    Err(e) => {
                        result.add_operation(false);
                        result.add_failed_resource(resource.resource_id.clone());
                        result.add_error(format!(
                            "Failed to rollback resource {}: {}",
                            resource.resource_id, e
                        ));
                        error!(
                            "Failed to rollback resource {}: {}",
                            resource.resource_id, e
                        );
                    }
                }
            } else {
                result.add_warning(format!(
                    "No rollback handler for provider: {}",
                    resource.provider
                ));
            }
        }

        Ok(())
    }

    /// Rollback a single task
    async fn rollback_single_task(&self, task: &WorkflowTask) -> Result<()> {
        // Update task status to cancelled
        self.storage
            .update_task_status(&task.id, TaskStatus::Cancelled)
            .await
            .context("Failed to update task status to cancelled")?;

        // Record audit entry for task rollback
        let audit_entry = AuditEntry {
            id: Uuid::new_v4().to_string(),
            task_id: task.id.clone(),
            operation: "task_rolled_back".to_string(),
            old_status: Some(task.status.clone()),
            new_status: Some(TaskStatus::Cancelled),
            user_id: None,
            timestamp: Utc::now(),
            metadata: {
                let mut meta = HashMap::new();
                meta.insert("task_name".to_string(), task.name.clone());
                meta
            },
        };

        self.storage
            .record_audit_entry(audit_entry)
            .await
            .context("Failed to record task rollback audit entry")?;

        Ok(())
    }

    /// Get active rollback operations
    pub async fn get_active_rollbacks(&self) -> HashSet<String> {
        self.active_rollbacks.lock().await.clone()
    }
}

/// Restores system state from checkpoints
pub struct StateRestorer {
    storage: Arc<dyn TaskStorage>,
    checkpoint_manager: Arc<CheckpointManager>,
    resource_tracker: Arc<ResourceTracker>,
}

impl StateRestorer {
    /// Create a new state restorer
    pub fn new(
        storage: Arc<dyn TaskStorage>,
        checkpoint_manager: Arc<CheckpointManager>,
        resource_tracker: Arc<ResourceTracker>,
    ) -> Self {
        Self {
            storage,
            checkpoint_manager,
            resource_tracker,
        }
    }

    /// Restore system state from a checkpoint
    pub async fn restore_from_checkpoint(&self, checkpoint_id: &str) -> Result<()> {
        let checkpoint = self
            .checkpoint_manager
            .get_checkpoint(checkpoint_id)
            .await
            .ok_or_else(|| anyhow::anyhow!("Checkpoint not found: {}", checkpoint_id))?;

        info!(
            "Restoring system state from checkpoint '{}' ({})",
            checkpoint.name, checkpoint_id
        );

        // Restore task states
        self.restore_task_states(&checkpoint)
            .await
            .context("Failed to restore task states")?;

        // Restore resource states
        self.restore_resource_states(&checkpoint)
            .await
            .context("Failed to restore resource states")?;

        info!(
            "Successfully restored system state from checkpoint '{}'",
            checkpoint.name
        );
        Ok(())
    }

    /// Restore task states from checkpoint
    async fn restore_task_states(&self, checkpoint: &Checkpoint) -> Result<()> {
        info!("Restoring {} task states", checkpoint.task_states.len());

        for task in &checkpoint.task_states {
            match self.storage.update_task(task.clone()).await {
                Ok(_) => debug!("Restored task state: {}", task.id),
                Err(e) => warn!("Failed to restore task state {}: {}", task.id, e),
            }
        }

        Ok(())
    }

    /// Restore resource states from checkpoint
    async fn restore_resource_states(&self, checkpoint: &Checkpoint) -> Result<()> {
        info!(
            "Restoring {} resource snapshots",
            checkpoint.resources.len()
        );

        for resource in &checkpoint.resources {
            match self.resource_tracker.track_resource(resource.clone()).await {
                Ok(_) => debug!("Restored resource state: {}", resource.resource_id),
                Err(e) => warn!(
                    "Failed to restore resource state {}: {}",
                    resource.resource_id, e
                ),
            }
        }

        Ok(())
    }
}

/// Main rollback system orchestrator
pub struct RollbackSystem {
    pub checkpoint_manager: Arc<CheckpointManager>,
    pub resource_tracker: Arc<ResourceTracker>,
    pub rollback_executor: Arc<RollbackExecutor>,
    pub state_restorer: Arc<StateRestorer>,
    config: RollbackConfig,
    auto_checkpoint_handle: Option<tokio::task::JoinHandle<()>>,
}

impl RollbackSystem {
    /// Create a new rollback system
    pub fn new(storage: Arc<dyn TaskStorage>, config: RollbackConfig) -> Self {
        let checkpoint_manager = Arc::new(CheckpointManager::new(storage.clone(), config.clone()));
        let resource_tracker = Arc::new(ResourceTracker::new(storage.clone()));
        let rollback_executor = Arc::new(RollbackExecutor::new(
            storage.clone(),
            resource_tracker.clone(),
            checkpoint_manager.clone(),
            config.clone(),
        ));
        let state_restorer = Arc::new(StateRestorer::new(
            storage.clone(),
            checkpoint_manager.clone(),
            resource_tracker.clone(),
        ));

        Self {
            checkpoint_manager,
            resource_tracker,
            rollback_executor,
            state_restorer,
            config,
            auto_checkpoint_handle: None,
        }
    }

    /// Initialize the rollback system
    pub async fn init(&mut self) -> Result<()> {
        info!(
            "Initializing rollback system with strategy: {:?}",
            self.config.strategy
        );

        // Start automatic checkpoint creation if enabled
        if self.config.auto_checkpoint_enabled {
            let checkpoint_manager = self.checkpoint_manager.clone();
            let interval =
                tokio::time::Duration::from_secs(self.config.checkpoint_interval_seconds);

            let handle = tokio::spawn(async move {
                let mut interval_timer = tokio::time::interval(interval);
                loop {
                    interval_timer.tick().await;

                    if checkpoint_manager.should_create_checkpoint().await {
                        match checkpoint_manager
                            .create_checkpoint(
                                format!("auto-{}", Utc::now().format("%Y%m%d-%H%M%S")),
                                Some("Automatic checkpoint".to_string()),
                            )
                            .await
                        {
                            Ok(id) => debug!("Created automatic checkpoint: {}", id),
                            Err(e) => error!("Failed to create automatic checkpoint: {}", e),
                        }
                    }
                }
            });

            self.auto_checkpoint_handle = Some(handle);
        }

        info!("Rollback system initialized successfully");
        Ok(())
    }

    /// Shutdown the rollback system
    pub async fn shutdown(&mut self) -> Result<()> {
        info!("Shutting down rollback system");

        if let Some(handle) = self.auto_checkpoint_handle.take() {
            handle.abort();
        }

        // Clean up expired checkpoints
        let removed = self
            .checkpoint_manager
            .cleanup_expired_checkpoints(30)
            .await
            .context("Failed to cleanup expired checkpoints")?;

        if removed > 0 {
            info!("Cleaned up {} expired checkpoints during shutdown", removed);
        }

        info!("Rollback system shutdown complete");
        Ok(())
    }

    /// Get system statistics
    pub async fn get_statistics(&self) -> RollbackStatistics {
        let checkpoints = self.checkpoint_manager.list_checkpoints().await;
        let tracked_resources = self.resource_tracker.get_tracked_resources().await;
        let active_rollbacks = self.rollback_executor.get_active_rollbacks().await;

        RollbackStatistics {
            total_checkpoints: checkpoints.len(),
            oldest_checkpoint: checkpoints.iter().map(|cp| cp.created_at).min(),
            newest_checkpoint: checkpoints.iter().map(|cp| cp.created_at).max(),
            tracked_resources: tracked_resources.len(),
            active_rollbacks: active_rollbacks.len(),
            auto_checkpoint_enabled: self.config.auto_checkpoint_enabled,
            checkpoint_interval_seconds: self.config.checkpoint_interval_seconds,
        }
    }
}

/// System statistics for monitoring
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RollbackStatistics {
    pub total_checkpoints: usize,
    pub oldest_checkpoint: Option<DateTime<Utc>>,
    pub newest_checkpoint: Option<DateTime<Utc>>,
    pub tracked_resources: usize,
    pub active_rollbacks: usize,
    pub auto_checkpoint_enabled: bool,
    pub checkpoint_interval_seconds: u64,
}

// Provider-specific rollback handlers for different cloud providers
pub mod providers {
    use super::*;

    /// UpCloud provider rollback handler
    pub struct UpCloudRollbackHandler {
        // UpCloud-specific client configuration would go here
    }

    #[async_trait]
    impl ProviderRollbackHandler for UpCloudRollbackHandler {
        fn provider_type(&self) -> &str {
            "upcloud"
        }

        async fn rollback_resource(&self, resource: &ResourceSnapshot) -> Result<()> {
            info!(
                "Rolling back UpCloud resource: {} ({})",
                resource.resource_id, resource.resource_type
            );

            match resource.resource_type.as_str() {
                "server" => {
                    // Implement server deletion logic
                    info!("Would delete UpCloud server: {}", resource.resource_id);
                }
                "storage" => {
                    // Implement storage deletion logic
                    info!("Would delete UpCloud storage: {}", resource.resource_id);
                }
                _ => {
                    warn!(
                        "Unknown UpCloud resource type for rollback: {}",
                        resource.resource_type
                    );
                }
            }

            Ok(())
        }

        async fn can_rollback_resource(&self, resource: &ResourceSnapshot) -> Result<bool> {
            // Check if resource can be safely rolled back
            Ok(matches!(
                resource.state,
                ResourceState::Active | ResourceState::Error { .. }
            ))
        }

        async fn get_resource_state(&self, _resource_id: &str) -> Result<ResourceState> {
            // Would query UpCloud API for current resource state
            Ok(ResourceState::Active)
        }

        async fn cleanup_resources_after(&self, checkpoint: &Checkpoint) -> Result<Vec<String>> {
            // Get list of UpCloud resources created after checkpoint
            let cleaned_resources = Vec::new();
            info!(
                "Would cleanup {} UpCloud resources created after checkpoint '{}'",
                cleaned_resources.len(),
                checkpoint.name
            );
            Ok(cleaned_resources)
        }

        async fn validate_rollback(&self, _checkpoint: &Checkpoint) -> Result<Vec<String>> {
            // Validate UpCloud-specific rollback conditions
            Ok(Vec::new()) // No validation errors
        }
    }

    /// AWS provider rollback handler
    pub struct AwsRollbackHandler {
        // AWS-specific client configuration would go here
    }

    #[async_trait]
    impl ProviderRollbackHandler for AwsRollbackHandler {
        fn provider_type(&self) -> &str {
            "aws"
        }

        async fn rollback_resource(&self, resource: &ResourceSnapshot) -> Result<()> {
            info!(
                "Rolling back AWS resource: {} ({})",
                resource.resource_id, resource.resource_type
            );

            match resource.resource_type.as_str() {
                "ec2_instance" => {
                    // Implement EC2 instance termination logic
                    info!("Would terminate AWS EC2 instance: {}", resource.resource_id);
                }
                "s3_bucket" => {
                    // Implement S3 bucket deletion logic
                    info!("Would delete AWS S3 bucket: {}", resource.resource_id);
                }
                "vpc" => {
                    // Implement VPC deletion logic
                    info!("Would delete AWS VPC: {}", resource.resource_id);
                }
                _ => {
                    warn!(
                        "Unknown AWS resource type for rollback: {}",
                        resource.resource_type
                    );
                }
            }

            Ok(())
        }

        async fn can_rollback_resource(&self, resource: &ResourceSnapshot) -> Result<bool> {
            // Check if resource can be safely rolled back
            Ok(matches!(
                resource.state,
                ResourceState::Active | ResourceState::Error { .. }
            ))
        }

        async fn get_resource_state(&self, _resource_id: &str) -> Result<ResourceState> {
            // Would query AWS API for current resource state
            Ok(ResourceState::Active)
        }

        async fn cleanup_resources_after(&self, checkpoint: &Checkpoint) -> Result<Vec<String>> {
            // Get list of AWS resources created after checkpoint
            let cleaned_resources = Vec::new();
            info!(
                "Would cleanup {} AWS resources created after checkpoint '{}'",
                cleaned_resources.len(),
                checkpoint.name
            );
            Ok(cleaned_resources)
        }

        async fn validate_rollback(&self, _checkpoint: &Checkpoint) -> Result<Vec<String>> {
            // Validate AWS-specific rollback conditions
            Ok(Vec::new()) // No validation errors
        }
    }

    /// Local provider rollback handler
    pub struct LocalRollbackHandler {
        // Local provider configuration
    }

    #[async_trait]
    impl ProviderRollbackHandler for LocalRollbackHandler {
        fn provider_type(&self) -> &str {
            "local"
        }

        async fn rollback_resource(&self, resource: &ResourceSnapshot) -> Result<()> {
            info!(
                "Rolling back local resource: {} ({})",
                resource.resource_id, resource.resource_type
            );

            match resource.resource_type.as_str() {
                "file" => {
                    // Implement file deletion logic
                    if let Some(path) = resource.config.get("path") {
                        info!("Would delete local file: {}", path);
                    }
                }
                "directory" => {
                    // Implement directory deletion logic
                    if let Some(path) = resource.config.get("path") {
                        info!("Would delete local directory: {}", path);
                    }
                }
                "process" => {
                    // Implement process termination logic
                    if let Some(pid) = resource.config.get("pid") {
                        info!("Would terminate local process: {}", pid);
                    }
                }
                _ => {
                    warn!(
                        "Unknown local resource type for rollback: {}",
                        resource.resource_type
                    );
                }
            }

            Ok(())
        }

        async fn can_rollback_resource(&self, resource: &ResourceSnapshot) -> Result<bool> {
            // Local resources can generally be rolled back
            Ok(!matches!(resource.state, ResourceState::Deleted))
        }

        async fn get_resource_state(&self, _resource_id: &str) -> Result<ResourceState> {
            // Check local resource state
            Ok(ResourceState::Active)
        }

        async fn cleanup_resources_after(&self, checkpoint: &Checkpoint) -> Result<Vec<String>> {
            // Get list of local resources created after checkpoint
            let cleaned_resources = Vec::new();
            info!(
                "Would cleanup {} local resources created after checkpoint '{}'",
                cleaned_resources.len(),
                checkpoint.name
            );
            Ok(cleaned_resources)
        }

        async fn validate_rollback(&self, _checkpoint: &Checkpoint) -> Result<Vec<String>> {
            // Validate local rollback conditions
            Ok(Vec::new()) // No validation errors
        }
    }
}

#[cfg(test)]
mod tests {
    use tempfile::TempDir;

    use super::*;
    use crate::storage::filesystem::FilesystemStorage;

    async fn create_test_storage() -> (Arc<dyn TaskStorage>, TempDir) {
        let temp_dir = TempDir::new().unwrap();
        let storage = FilesystemStorage::new(temp_dir.path().to_path_buf())
            .await
            .unwrap();
        (Arc::new(storage), temp_dir)
    }

    #[tokio::test]
    async fn test_checkpoint_creation() {
        let (storage, _temp_dir) = create_test_storage().await;
        let config = RollbackConfig::default();
        let checkpoint_manager = CheckpointManager::new(storage, config);

        let checkpoint_id = checkpoint_manager
            .create_checkpoint(
                "test_checkpoint".to_string(),
                Some("Test checkpoint".to_string()),
            )
            .await
            .unwrap();

        let checkpoint = checkpoint_manager.get_checkpoint(&checkpoint_id).await;
        assert!(checkpoint.is_some());
        assert_eq!(checkpoint.unwrap().name, "test_checkpoint");
    }

    #[tokio::test]
    async fn test_resource_tracking() {
        let (storage, _temp_dir) = create_test_storage().await;
        let resource_tracker = ResourceTracker::new(storage);

        let resource = ResourceSnapshot {
            resource_id: "test_resource".to_string(),
            resource_type: "server".to_string(),
            provider: "upcloud".to_string(),
            state: ResourceState::Active,
            created_at: Utc::now(),
            dependencies: Vec::new(),
            config: HashMap::new(),
            tags: HashMap::new(),
        };

        resource_tracker.track_resource(resource).await.unwrap();

        let tracked = resource_tracker.get_tracked_resources().await;
        assert_eq!(tracked.len(), 1);
        assert!(tracked.contains_key("test_resource"));
    }

    #[tokio::test]
    async fn test_rollback_system_initialization() {
        let (storage, _temp_dir) = create_test_storage().await;
        let config = RollbackConfig::default();
        let mut rollback_system = RollbackSystem::new(storage, config);

        rollback_system.init().await.unwrap();

        let stats = rollback_system.get_statistics().await;
        assert_eq!(stats.total_checkpoints, 0);
        assert_eq!(stats.tracked_resources, 0);
        assert_eq!(stats.active_rollbacks, 0);

        rollback_system.shutdown().await.unwrap();
    }
}
