//! Heterogeneous scheduler for multi-processor task assignment.
//!
//! Features:
//! - Optimal task-to-processor assignment
//! - Work stealing for load balancing
//! - Pipeline parallelism across processor types
//! - Dynamic rebalancing based on actual throughput

mod load_balancer;
mod work_queue;

pub use load_balancer::LoadBalancer;
pub use work_queue::WorkQueue;

use crate::device::DeviceRegistry;
use crate::error::ComputeError;
use crate::processor::{ProcessorId, ProcessorType};
use crate::task::{Task, TaskId};
use parking_lot::RwLock;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;

/// Heterogeneous scheduler that manages tasks across all processor types.
pub struct HeterogeneousScheduler {
    /// Device registry.
    device_registry: Arc<DeviceRegistry>,
    /// Per-processor-type task queues.
    queues: RwLock<HashMap<ProcessorType, WorkQueue>>,
    /// Load balancer.
    load_balancer: LoadBalancer,
    /// Active schedules.
    active_schedules: RwLock<HashMap<ScheduleId, Schedule>>,
}

impl HeterogeneousScheduler {
    /// Creates a new heterogeneous scheduler.
    pub fn new(device_registry: Arc<DeviceRegistry>) -> Self {
        Self {
            device_registry,
            queues: RwLock::new(HashMap::new()),
            load_balancer: LoadBalancer::new(),
            active_schedules: RwLock::new(HashMap::new()),
        }
    }

    /// Schedule a set of tasks for execution.
    pub async fn schedule(&self, tasks: Vec<Task>) -> Result<ScheduleResult, ComputeError> {
        if tasks.is_empty() {
            return Ok(ScheduleResult {
                schedule: Schedule::empty(),
                estimated_makespan: Duration::ZERO,
                processor_utilization: HashMap::new(),
            });
        }

        // 1. Build dependency graph
        let deps = self.build_dependency_graph(&tasks);

        // 2. Assign tasks to optimal processors
        let assignment = self.assign_tasks(&tasks, &deps).await?;

        // 3. Create execution schedule with stages
        let schedule = self.create_schedule(&tasks, &assignment, &deps)?;

        // 4. Estimate metrics
        let makespan = self.estimate_makespan(&schedule);
        let utilization = self.estimate_utilization(&schedule);

        // 5. Store active schedule
        self.active_schedules.write().insert(schedule.id, schedule.clone());

        Ok(ScheduleResult {
            schedule,
            estimated_makespan: makespan,
            processor_utilization: utilization,
        })
    }

    /// Execute a schedule.
    pub async fn execute(&self, schedule: &Schedule) -> Result<ExecutionResult, ComputeError> {
        let mut results = HashMap::new();
        let start = std::time::Instant::now();

        // Execute stages in order
        for stage in &schedule.stages {
            // Execute all tasks in this stage in parallel
            let mut handles = Vec::new();

            for task_id in &stage.tasks {
                let task = schedule.tasks.get(task_id)
                    .ok_or_else(|| ComputeError::Internal(format!("Task not found: {:?}", task_id)))?;
                let processor_id = schedule.assignment.get(task_id)
                    .ok_or_else(|| ComputeError::Internal(format!("No assignment for task: {:?}", task_id)))?;

                let processor = self.device_registry.get_processor(processor_id)?;
                let task_clone = task.clone();

                handles.push(tokio::spawn(async move {
                    processor.execute(task_clone.operation).await
                }));
            }

            // Wait for all tasks in stage
            for (i, handle) in handles.into_iter().enumerate() {
                let task_id = stage.tasks[i];
                match handle.await {
                    Ok(Ok(result)) => {
                        results.insert(task_id, TaskExecutionResult::Success(result));
                    }
                    Ok(Err(e)) => {
                        results.insert(task_id, TaskExecutionResult::Failed(e.to_string()));
                    }
                    Err(e) => {
                        results.insert(task_id, TaskExecutionResult::Failed(e.to_string()));
                    }
                }
            }
        }

        let total_time = start.elapsed();

        Ok(ExecutionResult {
            results,
            total_time,
            actual_utilization: self.measure_utilization(),
        })
    }

    /// Assign tasks to optimal processors.
    async fn assign_tasks(
        &self,
        tasks: &[Task],
        deps: &DependencyGraph,
    ) -> Result<TaskAssignment, ComputeError> {
        let mut assignment = TaskAssignment::new();

        // Sort tasks by priority and dependencies (topological sort)
        let sorted_tasks = self.topological_sort(tasks, deps);

        for task in sorted_tasks {
            // Find best processor for this task
            let best_processor = self.find_best_processor(&task).await?;

            // Check if we should rebalance
            let final_processor = self.load_balancer
                .maybe_rebalance(&task, best_processor, &assignment);

            assignment.assign(task.id, final_processor);
        }

        Ok(assignment)
    }

    /// Find the best processor for a task.
    async fn find_best_processor(&self, task: &Task) -> Result<ProcessorId, ComputeError> {
        let mut best_score = f64::NEG_INFINITY;
        let mut best_processor = None;

        // Get all available processors
        let processors = self.device_registry.all_processors();

        for processor in processors {
            if !processor.can_execute(&task.operation) {
                continue;
            }

            // Calculate score based on multiple factors
            let exec_time = processor.estimate_time(&task.operation);
            let energy = processor.estimate_energy(&task.operation);
            let load = processor.utilization();

            // Score = 1 / (time * (1 + load) * energy_factor)
            let time_factor = exec_time.as_secs_f64().max(0.001);
            let load_factor = 1.0 + load;
            let energy_factor = 1.0 + (energy / 1000.0); // Normalize energy

            let score = 1.0 / (time_factor * load_factor * energy_factor);

            if score > best_score {
                best_score = score;
                best_processor = Some(processor.id());
            }
        }

        best_processor.ok_or_else(|| {
            ComputeError::NoSuitableProcessor(format!("{:?}", task.operation.op_type()))
        })
    }

    /// Build dependency graph from tasks.
    fn build_dependency_graph(&self, tasks: &[Task]) -> DependencyGraph {
        let mut graph = DependencyGraph::new();

        for task in tasks {
            graph.add_node(task.id);
            for dep in &task.dependencies {
                graph.add_edge(*dep, task.id);
            }
        }

        graph
    }

    /// Topological sort of tasks respecting dependencies.
    fn topological_sort(&self, tasks: &[Task], deps: &DependencyGraph) -> Vec<Task> {
        let mut sorted = Vec::new();
        let mut visited = std::collections::HashSet::new();
        let task_map: HashMap<TaskId, Task> = tasks.iter()
            .map(|t| (t.id, t.clone()))
            .collect();

        fn visit(
            task_id: TaskId,
            task_map: &HashMap<TaskId, Task>,
            deps: &DependencyGraph,
            visited: &mut std::collections::HashSet<TaskId>,
            sorted: &mut Vec<Task>,
        ) {
            if visited.contains(&task_id) {
                return;
            }
            visited.insert(task_id);

            // Visit dependencies first
            if let Some(task_deps) = deps.dependencies.get(&task_id) {
                for dep in task_deps {
                    visit(*dep, task_map, deps, visited, sorted);
                }
            }

            if let Some(task) = task_map.get(&task_id) {
                sorted.push(task.clone());
            }
        }

        for task in tasks {
            visit(task.id, &task_map, deps, &mut visited, &mut sorted);
        }

        // Sort by priority within dependency constraints
        sorted.sort_by(|a, b| b.priority.cmp(&a.priority));

        sorted
    }

    /// Create execution schedule with parallel stages.
    fn create_schedule(
        &self,
        tasks: &[Task],
        assignment: &TaskAssignment,
        _deps: &DependencyGraph,
    ) -> Result<Schedule, ComputeError> {
        let mut stages = Vec::new();
        let mut scheduled = std::collections::HashSet::new();
        let task_map: HashMap<TaskId, Task> = tasks.iter()
            .map(|t| (t.id, t.clone()))
            .collect();

        while scheduled.len() < tasks.len() {
            let mut stage_tasks = Vec::new();

            for task in tasks {
                if scheduled.contains(&task.id) {
                    continue;
                }

                // Check if all dependencies are satisfied
                let deps_satisfied = task.dependencies.iter()
                    .all(|dep| scheduled.contains(dep));

                if deps_satisfied {
                    stage_tasks.push(task.id);
                }
            }

            if stage_tasks.is_empty() {
                return Err(ComputeError::SchedulingFailed(
                    "Circular dependency detected".to_string()
                ));
            }

            for task_id in &stage_tasks {
                scheduled.insert(*task_id);
            }

            stages.push(ScheduleStage {
                stage_id: stages.len(),
                tasks: stage_tasks,
            });
        }

        Ok(Schedule {
            id: ScheduleId::new(),
            tasks: task_map,
            assignment: assignment.clone(),
            stages,
        })
    }

    /// Estimate makespan (total execution time).
    fn estimate_makespan(&self, schedule: &Schedule) -> Duration {
        let mut total = Duration::ZERO;

        for stage in &schedule.stages {
            let mut max_stage_time = Duration::ZERO;

            for task_id in &stage.tasks {
                if let (Some(task), Some(proc_id)) = (
                    schedule.tasks.get(task_id),
                    schedule.assignment.get(task_id),
                ) {
                    if let Ok(processor) = self.device_registry.get_processor(proc_id) {
                        let time = processor.estimate_time(&task.operation);
                        max_stage_time = max_stage_time.max(time);
                    }
                }
            }

            total += max_stage_time;
        }

        total
    }

    /// Estimate processor utilization.
    fn estimate_utilization(&self, schedule: &Schedule) -> HashMap<ProcessorType, f64> {
        let mut work_time: HashMap<ProcessorType, Duration> = HashMap::new();
        let makespan = self.estimate_makespan(schedule);

        for task_id in schedule.assignment.assignments.keys() {
            if let (Some(task), Some(proc_id)) = (
                schedule.tasks.get(task_id),
                schedule.assignment.get(task_id),
            ) {
                if let Ok(processor) = self.device_registry.get_processor(proc_id) {
                    let proc_type = processor.processor_type();
                    let time = processor.estimate_time(&task.operation);
                    *work_time.entry(proc_type).or_default() += time;
                }
            }
        }

        work_time
            .into_iter()
            .map(|(proc_type, time)| {
                let utilization = if makespan.as_secs_f64() > 0.0 {
                    time.as_secs_f64() / makespan.as_secs_f64()
                } else {
                    0.0
                };
                (proc_type, utilization.min(1.0))
            })
            .collect()
    }

    /// Measure actual current utilization.
    fn measure_utilization(&self) -> HashMap<ProcessorType, f64> {
        let mut utilization = HashMap::new();

        for processor in self.device_registry.all_processors() {
            let proc_type = processor.processor_type();
            let util = processor.utilization();
            utilization
                .entry(proc_type)
                .and_modify(|u| *u = (*u + util) / 2.0)
                .or_insert(util);
        }

        utilization
    }
}

/// Schedule identifier.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct ScheduleId(pub u64);

impl ScheduleId {
    /// Creates a new schedule ID.
    pub fn new() -> Self {
        use rand::Rng;
        ScheduleId(rand::thread_rng().gen())
    }
}

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

/// Task-to-processor assignment.
#[derive(Clone, Debug, Default)]
pub struct TaskAssignment {
    /// Map from task ID to processor ID.
    pub assignments: HashMap<TaskId, ProcessorId>,
}

impl TaskAssignment {
    /// Creates a new empty assignment.
    pub fn new() -> Self {
        Self {
            assignments: HashMap::new(),
        }
    }

    /// Assigns a task to a processor.
    pub fn assign(&mut self, task_id: TaskId, processor_id: ProcessorId) {
        self.assignments.insert(task_id, processor_id);
    }

    /// Gets the assigned processor for a task.
    pub fn get(&self, task_id: &TaskId) -> Option<ProcessorId> {
        self.assignments.get(task_id).copied()
    }
}

/// Dependency graph for tasks.
#[derive(Clone, Debug, Default)]
pub struct DependencyGraph {
    /// Dependencies: task -> list of tasks it depends on.
    pub dependencies: HashMap<TaskId, Vec<TaskId>>,
    /// Dependents: task -> list of tasks that depend on it.
    pub dependents: HashMap<TaskId, Vec<TaskId>>,
}

impl DependencyGraph {
    /// Creates a new empty dependency graph.
    pub fn new() -> Self {
        Self {
            dependencies: HashMap::new(),
            dependents: HashMap::new(),
        }
    }

    /// Adds a node (task) to the graph.
    pub fn add_node(&mut self, task_id: TaskId) {
        self.dependencies.entry(task_id).or_default();
        self.dependents.entry(task_id).or_default();
    }

    /// Adds a dependency edge (from depends on to).
    pub fn add_edge(&mut self, from: TaskId, to: TaskId) {
        self.dependencies.entry(to).or_default().push(from);
        self.dependents.entry(from).or_default().push(to);
    }
}

/// Execution schedule.
#[derive(Clone, Debug)]
pub struct Schedule {
    /// Schedule ID.
    pub id: ScheduleId,
    /// All tasks.
    pub tasks: HashMap<TaskId, Task>,
    /// Task assignments.
    pub assignment: TaskAssignment,
    /// Execution stages (tasks within a stage can run in parallel).
    pub stages: Vec<ScheduleStage>,
}

impl Schedule {
    /// Creates an empty schedule.
    pub fn empty() -> Self {
        Self {
            id: ScheduleId::new(),
            tasks: HashMap::new(),
            assignment: TaskAssignment::new(),
            stages: Vec::new(),
        }
    }
}

/// A stage of parallel tasks.
#[derive(Clone, Debug)]
pub struct ScheduleStage {
    /// Stage index.
    pub stage_id: usize,
    /// Tasks in this stage (can run in parallel).
    pub tasks: Vec<TaskId>,
}

/// Result of scheduling.
#[derive(Clone, Debug)]
pub struct ScheduleResult {
    /// The schedule.
    pub schedule: Schedule,
    /// Estimated total execution time.
    pub estimated_makespan: Duration,
    /// Estimated processor utilization by type.
    pub processor_utilization: HashMap<ProcessorType, f64>,
}

/// Result of execution.
#[derive(Clone, Debug)]
pub struct ExecutionResult {
    /// Results per task.
    pub results: HashMap<TaskId, TaskExecutionResult>,
    /// Total execution time.
    pub total_time: Duration,
    /// Actual processor utilization.
    pub actual_utilization: HashMap<ProcessorType, f64>,
}

/// Result of a single task execution.
#[derive(Clone, Debug)]
pub enum TaskExecutionResult {
    /// Task completed successfully.
    Success(crate::processor::OperationResult),
    /// Task failed.
    Failed(String),
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::processor::Operation;
    use crate::task::{TaskPriority, TaskStatus};

    fn create_test_task(id: u64, op: Operation, deps: Vec<TaskId>) -> Task {
        Task {
            id: TaskId(id),
            operation: op,
            priority: TaskPriority::Normal,
            dependencies: deps,
            status: TaskStatus::Pending,
            deadline: None,
        }
    }

    #[test]
    fn test_dependency_graph() {
        let mut graph = DependencyGraph::new();

        graph.add_node(TaskId(1));
        graph.add_node(TaskId(2));
        graph.add_node(TaskId(3));

        graph.add_edge(TaskId(1), TaskId(2)); // 2 depends on 1
        graph.add_edge(TaskId(1), TaskId(3)); // 3 depends on 1
        graph.add_edge(TaskId(2), TaskId(3)); // 3 depends on 2

        assert_eq!(graph.dependencies[&TaskId(2)], vec![TaskId(1)]);
        assert_eq!(graph.dependencies[&TaskId(3)], vec![TaskId(1), TaskId(2)]);
    }

    #[test]
    fn test_task_assignment() {
        let mut assignment = TaskAssignment::new();

        assignment.assign(TaskId(1), ProcessorId(0));
        assignment.assign(TaskId(2), ProcessorId(1));

        assert_eq!(assignment.get(&TaskId(1)), Some(ProcessorId(0)));
        assert_eq!(assignment.get(&TaskId(2)), Some(ProcessorId(1)));
        assert_eq!(assignment.get(&TaskId(3)), None);
    }
}
