use std::collections::VecDeque;
use std::fmt;

use serde::{Deserialize, Serialize};

use arrdb_core::{ArrDbError, ChunkCoord, Result};
use arrdb_expr::AxisRange;

use crate::expr_ir::ExprProgram;
use crate::physical_plan::AggOp;

/// Unique task identifier within a task graph.
pub type TaskId = u32;

/// The kind of work a task performs.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum TaskKind {
    /// Read a chunk from an external data source.
    ReadChunk {
        array: String,
        coord: ChunkCoord,
        source_uri: String,
    },
    /// Apply a compiled expression program to a chunk.
    ApplyExpr { program: ExprProgram },
    /// Compute partial aggregation over a chunk.
    PartialAgg { aggs: Vec<AggOp> },
    /// Merge partial aggregation results.
    MergeAgg { aggs: Vec<AggOp> },
    /// Trim (spatial subsetting) a chunk.
    Trim { ranges: Vec<AxisRange> },
    /// Shuffle data to a target partition.
    Shuffle { target_partition: u32 },
    /// Emit the final result.
    Emit,
}

impl fmt::Display for TaskKind {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::ReadChunk { array, coord, .. } => write!(f, "ReadChunk({array}, {coord})"),
            Self::ApplyExpr { program } => {
                write!(f, "ApplyExpr(ops={})", program.ops.len())
            }
            Self::PartialAgg { aggs } => {
                let names: Vec<_> = aggs.iter().map(|a| a.to_string()).collect();
                write!(f, "PartialAgg({})", names.join(", "))
            }
            Self::MergeAgg { aggs } => {
                let names: Vec<_> = aggs.iter().map(|a| a.to_string()).collect();
                write!(f, "MergeAgg({})", names.join(", "))
            }
            Self::Trim { .. } => write!(f, "Trim"),
            Self::Shuffle { target_partition } => write!(f, "Shuffle(→{target_partition})"),
            Self::Emit => write!(f, "Emit"),
        }
    }
}

/// A single task in the scheduling DAG.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Task {
    /// Unique ID within the graph.
    pub id: TaskId,
    /// What this task does.
    pub kind: TaskKind,
    /// Which chunk this task operates on (if applicable).
    pub chunk_coord: Option<ChunkCoord>,
    /// Which partition/worker this task is assigned to.
    pub partition: u32,
}

/// A DAG of chunk-level tasks with explicit data dependencies.
///
/// Single-node: drives rayon scheduling.
/// Distributed (Phase 11): partitioned across workers.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct TaskGraph {
    tasks: Vec<Task>,
    /// Dependency edges: `(producer, consumer)`.
    edges: Vec<(TaskId, TaskId)>,
}

impl TaskGraph {
    /// Creates an empty task graph.
    pub fn new() -> Self {
        Self {
            tasks: Vec::new(),
            edges: Vec::new(),
        }
    }

    /// Adds a task to the graph, returning its ID.
    pub fn add_task(
        &mut self,
        kind: TaskKind,
        chunk_coord: Option<ChunkCoord>,
        partition: u32,
    ) -> TaskId {
        let id = self.tasks.len() as TaskId;
        self.tasks.push(Task {
            id,
            kind,
            chunk_coord,
            partition,
        });
        id
    }

    /// Adds a dependency edge: `consumer` depends on `producer`.
    pub fn add_edge(&mut self, producer: TaskId, consumer: TaskId) {
        self.edges.push((producer, consumer));
    }

    /// Returns the tasks that `task_id` depends on (its producers).
    pub fn dependencies(&self, task_id: TaskId) -> Vec<TaskId> {
        self.edges
            .iter()
            .filter(|(_, consumer)| *consumer == task_id)
            .map(|(producer, _)| *producer)
            .collect()
    }

    /// Number of tasks in the graph.
    pub fn num_tasks(&self) -> usize {
        self.tasks.len()
    }

    /// Returns a reference to a task by ID.
    pub fn task(&self, id: TaskId) -> Option<&Task> {
        self.tasks.get(id as usize)
    }

    /// Returns a topological ordering of task IDs (Kahn's algorithm).
    ///
    /// # Errors
    ///
    /// Returns an error if the graph contains a cycle.
    pub fn topological_order(&self) -> Result<Vec<TaskId>> {
        let n = self.tasks.len();
        let mut in_degree = vec![0u32; n];
        let mut adj: Vec<Vec<TaskId>> = vec![Vec::new(); n];

        for &(producer, consumer) in &self.edges {
            adj[producer as usize].push(consumer);
            in_degree[consumer as usize] += 1;
        }

        let mut queue: VecDeque<TaskId> = VecDeque::new();
        for (i, &deg) in in_degree.iter().enumerate() {
            if deg == 0 {
                queue.push_back(i as TaskId);
            }
        }

        let mut order = Vec::with_capacity(n);
        while let Some(task) = queue.pop_front() {
            order.push(task);
            for &next in &adj[task as usize] {
                in_degree[next as usize] -= 1;
                if in_degree[next as usize] == 0 {
                    queue.push_back(next);
                }
            }
        }

        if order.len() != n {
            return Err(ArrDbError::Expression("task graph contains a cycle".into()));
        }

        Ok(order)
    }
}

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

impl fmt::Display for TaskGraph {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        writeln!(
            f,
            "TaskGraph ({} tasks, {} edges):",
            self.tasks.len(),
            self.edges.len()
        )?;
        for task in &self.tasks {
            let deps = self.dependencies(task.id);
            if deps.is_empty() {
                writeln!(
                    f,
                    "  [{}] {} (partition={})",
                    task.id, task.kind, task.partition
                )?;
            } else {
                let dep_strs: Vec<_> = deps.iter().map(|d| d.to_string()).collect();
                writeln!(
                    f,
                    "  [{}] {} (partition={}, deps=[{}])",
                    task.id,
                    task.kind,
                    task.partition,
                    dep_strs.join(", ")
                )?;
            }
        }
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use arrdb_core::ScalarValue;
    use arrdb_expr::BinOp;

    use crate::expr_ir::{ExprOp, RegId};

    #[test]
    fn empty_graph() {
        let graph = TaskGraph::new();
        assert_eq!(graph.num_tasks(), 0);
        assert_eq!(graph.topological_order().unwrap(), Vec::<TaskId>::new());
    }

    #[test]
    fn linear_pipeline() {
        let mut graph = TaskGraph::new();

        let read = graph.add_task(
            TaskKind::ReadChunk {
                array: "climate".into(),
                coord: ChunkCoord::new(&[0, 0]),
                source_uri: "zarr:///data/climate.zarr".into(),
            },
            Some(ChunkCoord::new(&[0, 0])),
            0,
        );

        let apply = graph.add_task(
            TaskKind::ApplyExpr {
                program: ExprProgram {
                    ops: vec![ExprOp::LoadAttr {
                        dst: RegId(0),
                        attr_idx: 0,
                    }],
                    output_reg: RegId(0),
                    output_dtype: arrdb_core::DType::Float32,
                    num_registers: 1,
                },
            },
            Some(ChunkCoord::new(&[0, 0])),
            0,
        );

        let emit = graph.add_task(TaskKind::Emit, None, 0);

        graph.add_edge(read, apply);
        graph.add_edge(apply, emit);

        assert_eq!(graph.num_tasks(), 3);
        assert_eq!(graph.dependencies(read), Vec::<TaskId>::new());
        assert_eq!(graph.dependencies(apply), vec![read]);
        assert_eq!(graph.dependencies(emit), vec![apply]);

        let order = graph.topological_order().unwrap();
        assert_eq!(order, vec![0, 1, 2]);
    }

    #[test]
    fn diamond_dag() {
        // read_a and read_b → join → emit
        let mut graph = TaskGraph::new();

        let read_a = graph.add_task(
            TaskKind::ReadChunk {
                array: "a".into(),
                coord: ChunkCoord::new(&[0]),
                source_uri: "zarr:///a".into(),
            },
            Some(ChunkCoord::new(&[0])),
            0,
        );
        let read_b = graph.add_task(
            TaskKind::ReadChunk {
                array: "b".into(),
                coord: ChunkCoord::new(&[0]),
                source_uri: "zarr:///b".into(),
            },
            Some(ChunkCoord::new(&[0])),
            0,
        );
        let apply = graph.add_task(
            TaskKind::ApplyExpr {
                program: ExprProgram {
                    ops: vec![
                        ExprOp::LoadAttr {
                            dst: RegId(0),
                            attr_idx: 0,
                        },
                        ExprOp::LoadAttr {
                            dst: RegId(1),
                            attr_idx: 1,
                        },
                        ExprOp::BinaryOp {
                            op: BinOp::Sub,
                            dst: RegId(2),
                            lhs: RegId(0),
                            rhs: RegId(1),
                        },
                    ],
                    output_reg: RegId(2),
                    output_dtype: arrdb_core::DType::Float32,
                    num_registers: 3,
                },
            },
            Some(ChunkCoord::new(&[0])),
            0,
        );
        let emit = graph.add_task(TaskKind::Emit, None, 0);

        graph.add_edge(read_a, apply);
        graph.add_edge(read_b, apply);
        graph.add_edge(apply, emit);

        let order = graph.topological_order().unwrap();
        assert_eq!(order.len(), 4);

        // read_a and read_b must come before apply
        let pos_a = order.iter().position(|&x| x == read_a).unwrap();
        let pos_b = order.iter().position(|&x| x == read_b).unwrap();
        let pos_apply = order.iter().position(|&x| x == apply).unwrap();
        let pos_emit = order.iter().position(|&x| x == emit).unwrap();

        assert!(pos_a < pos_apply);
        assert!(pos_b < pos_apply);
        assert!(pos_apply < pos_emit);
    }

    #[test]
    fn cycle_detection() {
        let mut graph = TaskGraph::new();
        let t0 = graph.add_task(TaskKind::Emit, None, 0);
        let t1 = graph.add_task(TaskKind::Emit, None, 0);
        let t2 = graph.add_task(TaskKind::Emit, None, 0);

        graph.add_edge(t0, t1);
        graph.add_edge(t1, t2);
        graph.add_edge(t2, t0); // cycle!

        let err = graph.topological_order().unwrap_err();
        assert!(err.to_string().contains("cycle"));
    }

    #[test]
    fn task_graph_display() {
        let mut graph = TaskGraph::new();
        let read = graph.add_task(
            TaskKind::ReadChunk {
                array: "arr".into(),
                coord: ChunkCoord::new(&[0, 0]),
                source_uri: "zarr:///arr".into(),
            },
            Some(ChunkCoord::new(&[0, 0])),
            0,
        );
        let emit = graph.add_task(TaskKind::Emit, None, 0);
        graph.add_edge(read, emit);

        let display = graph.to_string();
        assert!(display.contains("TaskGraph (2 tasks, 1 edges)"));
        assert!(display.contains("ReadChunk"));
        assert!(display.contains("Emit"));
        assert!(display.contains("deps=[0]"));
    }

    #[test]
    fn task_kind_display() {
        assert_eq!(TaskKind::Emit.to_string(), "Emit");
        assert_eq!(
            TaskKind::Shuffle {
                target_partition: 3
            }
            .to_string(),
            "Shuffle(→3)"
        );
        assert_eq!(
            TaskKind::ReadChunk {
                array: "a".into(),
                coord: ChunkCoord::new(&[1, 2]),
                source_uri: "zarr:///a".into()
            }
            .to_string(),
            "ReadChunk(a, (1, 2))"
        );
    }

    #[test]
    fn task_graph_serde_roundtrip() {
        let mut graph = TaskGraph::new();
        let t0 = graph.add_task(
            TaskKind::ReadChunk {
                array: "x".into(),
                coord: ChunkCoord::new(&[0]),
                source_uri: "zarr:///x".into(),
            },
            Some(ChunkCoord::new(&[0])),
            0,
        );
        let t1 = graph.add_task(TaskKind::Emit, None, 0);
        graph.add_edge(t0, t1);

        let json = serde_json::to_string(&graph).unwrap();
        let back: TaskGraph = serde_json::from_str(&json).unwrap();
        assert_eq!(back, graph);
    }

    #[test]
    fn partial_agg_display() {
        let kind = TaskKind::PartialAgg {
            aggs: vec![
                crate::physical_plan::AggOp {
                    func: crate::physical_plan::AggFunc::Sum,
                    attr: "temp".into(),
                    output_dtype: arrdb_core::DType::Float64,
                },
                crate::physical_plan::AggOp {
                    func: crate::physical_plan::AggFunc::Count,
                    attr: "temp".into(),
                    output_dtype: arrdb_core::DType::Int64,
                },
            ],
        };
        assert_eq!(kind.to_string(), "PartialAgg(SUM(temp), COUNT(temp))");
    }

    #[test]
    fn trim_kind() {
        let kind = TaskKind::Trim {
            ranges: vec![AxisRange::Full, AxisRange::Trim(10, 50)],
        };
        assert_eq!(kind.to_string(), "Trim");
    }

    #[test]
    fn apply_expr_display() {
        let kind = TaskKind::ApplyExpr {
            program: ExprProgram {
                ops: vec![
                    ExprOp::LoadAttr {
                        dst: RegId(0),
                        attr_idx: 0,
                    },
                    ExprOp::ScalarBinaryOp {
                        op: BinOp::Mul,
                        dst: RegId(1),
                        lhs: RegId(0),
                        rhs: ScalarValue::Float64(1.8),
                    },
                ],
                output_reg: RegId(1),
                output_dtype: arrdb_core::DType::Float64,
                num_registers: 2,
            },
        };
        assert_eq!(kind.to_string(), "ApplyExpr(ops=2)");
    }
}
