//! DAG (Directed Acyclic Graph) scheduling system with Airflow-like capabilities
//!
//! Provides comprehensive task orchestration including:
//! - Task dependency resolution with topological sorting
//! - Parallel execution of independent tasks
//! - Retry mechanisms with exponential backoff
//! - Conditional task execution based on upstream results
//! - Dynamic DAG generation and modification
//! - Resource allocation and priority scheduling

use chrono::{DateTime, Utc};
use petgraph::algo::{is_cyclic_directed, toposort};
use petgraph::{Directed, Graph};
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tokio::sync::RwLock;
use tracing::{debug, instrument, warn};
use uuid::Uuid;

use crate::{
    EtlError, Result,
    connectors::{DataSink, DataSource},
    monitoring::EtlMetrics,
    state::StateManager,
    transform::Transformation,
};

/// Configuration for DAG scheduling
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DagConfig {
    /// Maximum number of task retries
    pub max_retries: u32,
    /// Base retry delay in milliseconds
    pub base_retry_delay_ms: u64,
    /// Maximum retry delay in milliseconds
    pub max_retry_delay_ms: u64,
    /// Task execution timeout in milliseconds
    pub task_timeout_ms: u64,
    /// Enable dynamic DAG modifications
    pub allow_dynamic_dag: bool,
    /// Maximum DAG depth
    pub max_dag_depth: usize,
    /// Enable priority-based scheduling
    pub priority_scheduling: bool,
}

impl Default for DagConfig {
    fn default() -> Self {
        Self {
            max_retries: 3,
            base_retry_delay_ms: 1000,
            max_retry_delay_ms: 60000,
            task_timeout_ms: 300000, // 5 minutes
            allow_dynamic_dag: true,
            max_dag_depth: 50,
            priority_scheduling: true,
        }
    }
}

/// DAG scheduler responsible for task orchestration
pub struct DagScheduler {
    /// Configuration
    config: DagConfig,
    /// State manager for persistence
    state_manager: Arc<StateManager>,
    /// Metrics collector
    metrics: Arc<EtlMetrics>,
    /// Active DAG executions
    active_executions: Arc<RwLock<HashMap<Uuid, DagExecution>>>,
}

/// Represents a DAG of tasks with dependencies
#[derive(Debug, Clone)]
pub struct TaskGraph {
    /// Tasks in the graph
    pub tasks: HashMap<String, Task>,
    /// Task dependencies (`task_id` -> list of dependencies)
    pub dependencies: HashMap<String, Vec<String>>,
    /// Graph metadata
    pub metadata: GraphMetadata,
}

/// Metadata for a task graph
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GraphMetadata {
    /// Graph identifier
    pub graph_id: String,
    /// Graph name
    pub name: String,
    /// Description
    pub description: Option<String>,
    /// Tags for categorization
    pub tags: Vec<String>,
    /// Creation timestamp
    pub created_at: DateTime<Utc>,
    /// Last modified timestamp
    pub modified_at: DateTime<Utc>,
    /// Version
    pub version: String,
}

/// Individual task in a DAG
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Task {
    /// Task identifier
    pub id: String,
    /// Task name
    pub name: String,
    /// Task description
    pub description: Option<String>,
    /// Task type
    pub task_type: TaskType,
    /// Data source configuration
    pub source: Option<DataSource>,
    /// Data transformation configuration
    pub transformation: Option<Transformation>,
    /// Data sink configuration
    pub sink: Option<DataSink>,
    /// Task dependencies
    pub dependencies: Vec<String>,
    /// Task priority (higher number = higher priority)
    pub priority: i32,
    /// Retry configuration
    pub retry_config: RetryConfig,
    /// Resource requirements
    pub resources: ResourceRequirements,
    /// Conditional execution configuration
    pub condition: Option<TaskCondition>,
    /// Task timeout override
    pub timeout_ms: Option<u64>,
    /// Task metadata
    pub metadata: HashMap<String, String>,
    /// Task configuration parameters
    pub config: HashMap<String, serde_json::Value>,
}

/// Task priority levels
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord, Default)]
pub enum TaskPriority {
    Low = 1,
    #[default]
    Medium = 2,
    High = 3,
    Critical = 4,
}

/// Types of tasks that can be executed
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum TaskType {
    /// Extract data from a source
    Extract,
    /// Transform data using specified transformations
    Transform,
    /// Load data to a destination
    Load,
    /// Combined ETL operation
    Etl,
    /// Data quality check
    QualityCheck,
    /// Custom task with user-defined logic
    Custom(String),
    /// Conditional branching task
    Branch,
    /// Wait/delay task
    Wait(Duration),
}

impl std::fmt::Display for TaskType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Extract => write!(f, "Extract"),
            Self::Transform => write!(f, "Transform"),
            Self::Load => write!(f, "Load"),
            Self::Etl => write!(f, "ETL"),
            Self::QualityCheck => write!(f, "QualityCheck"),
            Self::Custom(name) => write!(f, "Custom({name})"),
            Self::Branch => write!(f, "Branch"),
            Self::Wait(duration) => write!(f, "Wait({duration:?})"),
        }
    }
}

impl TaskType {
    #[must_use]
    pub fn as_str(&self) -> &str {
        match self {
            Self::Extract => "extract",
            Self::Transform => "transform",
            Self::Load => "load",
            Self::Etl => "etl",
            Self::QualityCheck => "quality_check",
            Self::Custom(_) => "custom",
            Self::Branch => "branch",
            Self::Wait(_) => "wait",
        }
    }

    pub fn from_str(s: &str) -> Result<Self> {
        match s.to_lowercase().as_str() {
            "extract" => Ok(Self::Extract),
            "transform" => Ok(Self::Transform),
            "load" => Ok(Self::Load),
            "etl" => Ok(Self::Etl),
            "quality_check" => Ok(Self::QualityCheck),
            "branch" => Ok(Self::Branch),
            _ if s.starts_with("custom:") => {
                let name = s.strip_prefix("custom:").unwrap_or(s);
                Ok(Self::Custom(name.to_string()))
            }
            _ => Err(EtlError::validation(format!("Unknown task type: {s}"))),
        }
    }
}

impl Task {
    #[must_use]
    pub fn new(id: String, task_type: TaskType) -> Self {
        Self {
            id,
            name: String::new(),
            description: None,
            task_type,
            source: None,
            transformation: None,
            sink: None,
            dependencies: Vec::new(),
            retry_config: RetryConfig::default(),
            priority: 0,
            resources: ResourceRequirements::default(),
            condition: None,
            timeout_ms: None,
            metadata: HashMap::new(),
            config: HashMap::new(),
        }
    }
}

/// Retry configuration for tasks
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RetryConfig {
    /// Maximum number of retries
    pub max_retries: u32,
    /// Retry strategy
    pub strategy: RetryStrategy,
    /// Base delay between retries
    pub base_delay_ms: u64,
    /// Maximum delay between retries
    pub max_delay_ms: u64,
    /// Exponential backoff multiplier
    pub backoff_multiplier: f64,
}

/// Retry strategies for failed tasks
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum RetryStrategy {
    /// Fixed delay between retries
    Fixed,
    /// Exponential backoff with jitter
    ExponentialBackoff,
    /// Linear backoff
    LinearBackoff,
    /// Custom retry logic
    Custom(String),
}

/// Resource requirements for task execution
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ResourceRequirements {
    /// CPU cores required
    pub cpu_cores: f32,
    /// Memory required in MB
    pub memory_mb: u64,
    /// Disk space required in MB
    pub disk_mb: u64,
    /// GPU resources required
    pub gpu_memory_mb: Option<u64>,
    /// Network bandwidth required in Mbps
    pub network_mbps: Option<u64>,
}

/// Condition for task execution
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskCondition {
    /// Condition type
    pub condition_type: ConditionType,
    /// Condition expression
    pub expression: String,
    /// Task to execute if condition is true
    pub on_true: Option<String>,
    /// Task to execute if condition is false
    pub on_false: Option<String>,
}

/// Types of task conditions
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ConditionType {
    /// Condition based on upstream task status
    UpstreamStatus,
    /// Condition based on data availability
    DataAvailability,
    /// Custom condition logic
    Custom(String),
    /// Time-based condition
    TimeBased,
    /// Resource availability condition
    ResourceAvailability,
}

/// Status of a task execution
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub enum TaskStatus {
    /// Task is pending execution
    Pending,
    /// Task is currently running
    Running,
    /// Task completed successfully
    Success,
    /// Task failed
    Failed,
    /// Task was skipped due to conditions
    Skipped,
    /// Task was cancelled
    Cancelled,
    /// Task timed out
    TimedOut,
    /// Task is waiting for retry
    Retrying,
}

/// Result of task execution
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskResult {
    /// Task identifier
    pub task_id: String,
    /// Execution status
    pub status: TaskStatus,
    /// Start time
    pub start_time: DateTime<Utc>,
    /// End time
    pub end_time: Option<DateTime<Utc>>,
    /// Records processed
    pub records_processed: u64,
    /// Bytes processed
    pub bytes_processed: u64,
    /// Error message if failed
    pub error_message: Option<String>,
    /// Output data location
    pub output_location: Option<String>,
    /// Task metrics
    pub metrics: TaskMetrics,
}

/// Metrics for task execution
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TaskMetrics {
    /// Execution duration in milliseconds
    pub duration_ms: u64,
    /// Peak memory usage in bytes
    pub peak_memory_bytes: u64,
    /// CPU time used in milliseconds
    pub cpu_time_ms: u64,
    /// I/O operations count
    pub io_operations: u64,
    /// Network bytes transferred
    pub network_bytes: u64,
}

/// Execution context for a DAG
#[derive(Debug)]
pub struct DagExecution {
    /// Execution identifier
    pub execution_id: Uuid,
    /// Task graph being executed
    pub graph: TaskGraph,
    /// Execution plan
    pub plan: ExecutionPlan,
    /// Task statuses
    pub task_status: HashMap<String, TaskStatus>,
    /// Task results
    pub task_results: HashMap<String, TaskResult>,
    /// Execution start time
    pub start_time: Instant,
    /// Current execution stage
    pub current_stage: usize,
}

/// Execution plan for a DAG
#[derive(Debug, Clone)]
pub struct ExecutionPlan {
    /// Execution stages (tasks that can run in parallel)
    pub stages: Vec<ExecutionStage>,
    /// Total estimated duration
    pub estimated_duration: Duration,
    /// Resource requirements per stage
    pub stage_resources: Vec<ResourceRequirements>,
}

/// A stage of execution containing tasks that can run in parallel
#[derive(Debug, Clone)]
pub struct ExecutionStage {
    /// Stage identifier
    pub stage_id: usize,
    /// Tasks in this stage
    pub tasks: Vec<Task>,
    /// Stage dependencies (previous stages that must complete)
    pub dependencies: Vec<usize>,
}

/// Context for task execution
#[derive(Debug)]
pub struct ExecutionContext {
    /// Task being executed
    pub task: Task,
    /// Execution start time
    pub start_time: Instant,
    /// Retry attempt number
    pub retry_attempt: u32,
    /// Execution metadata
    pub metadata: HashMap<String, String>,
}

/// Builder for creating tasks
pub struct TaskBuilder {
    task: Task,
}

impl DagScheduler {
    /// Create a new DAG scheduler
    pub async fn new(
        config: DagConfig,
        state_manager: Arc<StateManager>,
        metrics: Arc<EtlMetrics>,
    ) -> Result<Self> {
        Ok(Self {
            config,
            state_manager,
            metrics,
            active_executions: Arc::new(RwLock::new(HashMap::new())),
        })
    }

    /// Plan the execution of a DAG
    #[instrument(skip(self, graph))]
    pub async fn plan_execution(&self, graph: &TaskGraph) -> Result<ExecutionPlan> {
        debug!(
            "Planning DAG execution for graph: {}",
            graph.metadata.graph_id
        );

        // Validate the graph
        graph.validate()?;

        // Create execution stages using topological sort
        let stages = self.create_execution_stages(graph)?;

        // Estimate execution duration
        let estimated_duration = self.estimate_execution_duration(&stages);

        // Calculate resource requirements
        let stage_resources = self.calculate_stage_resources(&stages);

        Ok(ExecutionPlan {
            stages,
            estimated_duration,
            stage_resources,
        })
    }

    /// Create execution stages from task graph
    fn create_execution_stages(&self, graph: &TaskGraph) -> Result<Vec<ExecutionStage>> {
        // Build petgraph from task dependencies
        let mut pg = Graph::<String, (), Directed>::new();
        let mut node_indices = HashMap::new();

        // Add nodes
        for task_id in graph.tasks.keys() {
            let idx = pg.add_node(task_id.clone());
            node_indices.insert(task_id.clone(), idx);
        }

        // Add edges for dependencies
        for (task_id, deps) in &graph.dependencies {
            let task_idx = node_indices[task_id];
            for dep in deps {
                if let Some(&dep_idx) = node_indices.get(dep) {
                    pg.add_edge(dep_idx, task_idx, ());
                }
            }
        }

        // Check for cycles
        if is_cyclic_directed(&pg) {
            return Err(EtlError::DagValidation("Graph contains cycles".to_string()));
        }

        // Perform topological sort
        let topo_sort = toposort(&pg, None).map_err(|_| {
            EtlError::DagValidation("Failed to sort graph topologically".to_string())
        })?;

        // Group tasks into stages based on dependency levels
        let mut stages = Vec::new();
        let mut task_levels = HashMap::new();

        // Calculate dependency levels
        for node_idx in topo_sort {
            let task_id = &pg[node_idx];
            let task = &graph.tasks[task_id];

            let level = if task.dependencies.is_empty() {
                0
            } else {
                task.dependencies
                    .iter()
                    .map(|dep| task_levels.get(dep).unwrap_or(&0) + 1)
                    .max()
                    .unwrap_or(0)
            };

            task_levels.insert(task_id.clone(), level);
        }

        // Group tasks by level into stages
        let max_level = task_levels.values().max().unwrap_or(&0);
        for stage_id in 0..=*max_level {
            let stage_tasks: Vec<Task> = task_levels
                .iter()
                .filter(|(_, level)| **level == stage_id)
                .map(|(task_id, _)| graph.tasks[task_id].clone())
                .collect();

            if !stage_tasks.is_empty() {
                stages.push(ExecutionStage {
                    stage_id,
                    tasks: stage_tasks,
                    dependencies: if stage_id > 0 {
                        vec![stage_id - 1]
                    } else {
                        vec![]
                    },
                });
            }
        }

        Ok(stages)
    }

    /// Estimate total execution duration
    fn estimate_execution_duration(&self, stages: &[ExecutionStage]) -> Duration {
        let mut total_duration = Duration::ZERO;

        for stage in stages {
            // For parallel execution, use the longest task in the stage
            let stage_duration = stage
                .tasks
                .iter()
                .map(|task| {
                    Duration::from_millis(task.timeout_ms.unwrap_or(self.config.task_timeout_ms))
                })
                .max()
                .unwrap_or(Duration::ZERO);

            total_duration += stage_duration;
        }

        total_duration
    }

    /// Calculate resource requirements for each stage
    fn calculate_stage_resources(&self, stages: &[ExecutionStage]) -> Vec<ResourceRequirements> {
        stages
            .iter()
            .map(|stage| {
                // Sum resources for all tasks in the stage (parallel execution)
                let total_cpu: f32 = stage.tasks.iter().map(|t| t.resources.cpu_cores).sum();
                let total_memory: u64 = stage.tasks.iter().map(|t| t.resources.memory_mb).sum();
                let total_disk: u64 = stage.tasks.iter().map(|t| t.resources.disk_mb).sum();

                ResourceRequirements {
                    cpu_cores: total_cpu,
                    memory_mb: total_memory,
                    disk_mb: total_disk,
                    gpu_memory_mb: None,
                    network_mbps: None,
                }
            })
            .collect()
    }
}

impl TaskGraph {
    /// Create a new empty task graph
    #[must_use]
    pub fn new() -> Self {
        Self {
            tasks: HashMap::new(),
            dependencies: HashMap::new(),
            metadata: GraphMetadata {
                graph_id: Uuid::new_v4().to_string(),
                name: "unnamed_graph".to_string(),
                description: None,
                tags: Vec::new(),
                created_at: Utc::now(),
                modified_at: Utc::now(),
                version: "1.0.0".to_string(),
            },
        }
    }

    /// Add a task to the graph
    pub fn add_task(&mut self, task: Task) {
        let deps = task.dependencies.clone();
        self.dependencies.insert(task.id.clone(), deps);
        self.tasks.insert(task.id.clone(), task);
        self.metadata.modified_at = Utc::now();
    }

    /// Remove a task from the graph
    pub fn remove_task(&mut self, task_id: &str) -> Option<Task> {
        self.dependencies.remove(task_id);

        // Remove this task from other tasks' dependencies
        for deps in self.dependencies.values_mut() {
            deps.retain(|dep| dep != task_id);
        }

        self.metadata.modified_at = Utc::now();
        self.tasks.remove(task_id)
    }

    /// Add a dependency between tasks
    pub fn add_dependency(&mut self, task_id: &str, dependency: &str) -> Result<()> {
        if !self.tasks.contains_key(task_id) {
            return Err(EtlError::DagValidation(format!("Task {task_id} not found")));
        }

        if !self.tasks.contains_key(dependency) {
            return Err(EtlError::DagValidation(format!(
                "Dependency task {dependency} not found"
            )));
        }

        self.dependencies
            .entry(task_id.to_string())
            .or_default()
            .push(dependency.to_string());

        // Update the task's dependencies as well
        if let Some(task) = self.tasks.get_mut(task_id)
            && !task.dependencies.contains(&dependency.to_string())
        {
            task.dependencies.push(dependency.to_string());
        }

        self.metadata.modified_at = Utc::now();

        // Check for cycles after adding dependency
        self.validate_no_cycles()
    }

    /// Validate the graph for correctness
    pub fn validate(&self) -> Result<()> {
        // Check for missing dependencies
        for (task_id, deps) in &self.dependencies {
            for dep in deps {
                if !self.tasks.contains_key(dep) {
                    return Err(EtlError::DagValidation(format!(
                        "Task {task_id} depends on non-existent task {dep}"
                    )));
                }
            }
        }

        // Check for cycles
        self.validate_no_cycles()?;

        // Check for orphaned tasks (tasks with no dependencies and no dependents)
        let mut has_dependents = HashSet::new();
        for deps in self.dependencies.values() {
            for dep in deps {
                has_dependents.insert(dep.clone());
            }
        }

        for task_id in self.tasks.keys() {
            if self
                .dependencies
                .get(task_id)
                .is_none_or(std::vec::Vec::is_empty)
                && !has_dependents.contains(task_id)
            {
                warn!(
                    "Task {} is orphaned (no dependencies or dependents)",
                    task_id
                );
            }
        }

        Ok(())
    }

    /// Validate that the graph contains no cycles
    fn validate_no_cycles(&self) -> Result<()> {
        let mut visited = HashSet::new();
        let mut rec_stack = HashSet::new();

        for task_id in self.tasks.keys() {
            if !visited.contains(task_id)
                && self.has_cycle_dfs(task_id, &mut visited, &mut rec_stack)
            {
                return Err(EtlError::DagValidation("Graph contains cycles".to_string()));
            }
        }

        Ok(())
    }

    /// DFS cycle detection helper
    fn has_cycle_dfs(
        &self,
        task_id: &str,
        visited: &mut HashSet<String>,
        rec_stack: &mut HashSet<String>,
    ) -> bool {
        visited.insert(task_id.to_string());
        rec_stack.insert(task_id.to_string());

        if let Some(deps) = self.dependencies.get(task_id) {
            for dep in deps {
                if !visited.contains(dep) {
                    if self.has_cycle_dfs(dep, visited, rec_stack) {
                        return true;
                    }
                } else if rec_stack.contains(dep) {
                    return true;
                }
            }
        }

        rec_stack.remove(task_id);
        false
    }
}

impl TaskBuilder {
    /// Create a new task builder
    #[must_use]
    pub fn new(id: &str) -> Self {
        Self {
            task: Task::new(id.to_string(), TaskType::Transform),
        }
    }

    /// Set task name
    #[must_use]
    pub fn name(mut self, name: &str) -> Self {
        self.task.name = name.to_string();
        self
    }

    /// Set task description
    #[must_use]
    pub fn description(mut self, description: &str) -> Self {
        self.task.description = Some(description.to_string());
        self
    }

    /// Set task type
    #[must_use]
    pub fn task_type(mut self, task_type: TaskType) -> Self {
        self.task.task_type = task_type;
        self
    }

    /// Add a dependency
    #[must_use]
    pub fn depends_on(mut self, dependency: &str) -> Self {
        self.task.dependencies.push(dependency.to_string());
        self
    }

    /// Set data source
    #[must_use]
    pub fn with_source(mut self, source: DataSource) -> Self {
        self.task.source = Some(source);
        self
    }

    /// Set transformation
    #[must_use]
    pub fn with_transformation(mut self, transformation: Transformation) -> Self {
        self.task.transformation = Some(transformation);
        self
    }

    /// Set data sink
    #[must_use]
    pub fn with_sink(mut self, sink: DataSink) -> Self {
        self.task.sink = Some(sink);
        self
    }

    /// Set task priority
    #[must_use]
    pub fn priority(mut self, priority: i32) -> Self {
        self.task.priority = priority;
        self
    }

    /// Set retry configuration
    #[must_use]
    pub fn retry_config(mut self, config: RetryConfig) -> Self {
        self.task.retry_config = config;
        self
    }

    /// Set resource requirements
    #[must_use]
    pub fn resources(mut self, resources: ResourceRequirements) -> Self {
        self.task.resources = resources;
        self
    }

    /// Build the task
    #[must_use]
    pub fn build(self) -> Task {
        self.task
    }
}

impl Default for RetryConfig {
    fn default() -> Self {
        Self {
            max_retries: 3,
            strategy: RetryStrategy::ExponentialBackoff,
            base_delay_ms: 1000,
            max_delay_ms: 60000,
            backoff_multiplier: 2.0,
        }
    }
}

impl Default for ResourceRequirements {
    fn default() -> Self {
        Self {
            cpu_cores: 1.0,
            memory_mb: 512,
            disk_mb: 1024,
            gpu_memory_mb: None,
            network_mbps: None,
        }
    }
}

impl Default for TaskGraph {
    fn default() -> Self {
        Self::new()
    }
}