use std::fmt;

use serde::{Deserialize, Serialize};

use arrdb_core::{ArraySchema, ChunkCoord, DType};
use arrdb_expr::AxisRange;

use crate::expr_ir::ExprProgram;

/// Aggregate function identifier used in partial/final aggregation nodes.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum AggFunc {
    Sum,
    Avg,
    Min,
    Max,
    Count,
}

impl fmt::Display for AggFunc {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Sum => write!(f, "SUM"),
            Self::Avg => write!(f, "AVG"),
            Self::Min => write!(f, "MIN"),
            Self::Max => write!(f, "MAX"),
            Self::Count => write!(f, "COUNT"),
        }
    }
}

/// A single aggregation operation — function + target attribute.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct AggOp {
    pub func: AggFunc,
    pub attr: String,
    pub output_dtype: DType,
}

impl fmt::Display for AggOp {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "{}({})", self.func, self.attr)
    }
}

/// Estimated statistics for a plan node (used for cost-based decisions).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct PlanStats {
    /// Estimated number of output chunks.
    pub est_chunks: u64,
    /// Estimated bytes per chunk.
    pub est_bytes_per_chunk: u64,
}

/// How to distribute data across threads/nodes.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum Partitioning {
    /// Round-robin across N partitions.
    RoundRobin(usize),
    /// Hash chunk coordinate into N buckets.
    ChunkHash(usize),
    /// Range-partition along a dimension.
    DimensionRange {
        dim_name: String,
        num_partitions: usize,
    },
}

impl Partitioning {
    /// Number of partitions.
    pub fn num_partitions(&self) -> usize {
        match self {
            Self::RoundRobin(n) | Self::ChunkHash(n) => *n,
            Self::DimensionRange { num_partitions, .. } => *num_partitions,
        }
    }
}

impl fmt::Display for Partitioning {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::RoundRobin(n) => write!(f, "RoundRobin({n})"),
            Self::ChunkHash(n) => write!(f, "ChunkHash({n})"),
            Self::DimensionRange {
                dim_name,
                num_partitions,
            } => write!(f, "DimensionRange({dim_name}, {num_partitions})"),
        }
    }
}

/// A serializable physical execution plan — the "what to execute" IR.
///
/// Produced by the physical planner (Phase 7), consumed by the executor (Phase 8).
/// Tree-structured: each node has zero or more child plans.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum PhysicalPlan {
    /// Read chunks from an external data source.
    Scan {
        array: String,
        schema: ArraySchema,
        source_uri: String,
        chunk_list: Vec<ChunkCoord>,
        projection: Vec<usize>,
        use_cache: bool,
        stats: Option<PlanStats>,
    },
    /// Trim (spatial subsetting) along one or more axes.
    Trim {
        input: Box<PhysicalPlan>,
        /// Chunk-local trim ranges.
        ranges: Vec<AxisRange>,
        /// Original global trim ranges (before chunk-local conversion).
        /// Used by `extract_dim_mapping` to compute correct offsets into
        /// the coordinate index. `None` when no localization was needed.
        global_ranges: Option<Vec<AxisRange>>,
    },
    /// Slice a single dimension at a fixed index.
    Slice {
        input: Box<PhysicalPlan>,
        axis: usize,
        /// Chunk-local index.
        index: i64,
        /// Original global index (before chunk-local conversion).
        /// Used by `extract_dim_mapping` for coordinate label lookup.
        global_index: Option<i64>,
    },
    /// Apply a compiled expression program, producing a new attribute.
    Apply {
        input: Box<PhysicalPlan>,
        program: ExprProgram,
        output_attr: String,
    },
    /// Filter chunks using a boolean predicate program.
    Filter {
        input: Box<PhysicalPlan>,
        predicate: ExprProgram,
    },
    /// Per-chunk partial aggregation.
    PartialAggregate {
        input: Box<PhysicalPlan>,
        aggs: Vec<AggOp>,
    },
    /// Merge partial aggregation results.
    FinalAggregate {
        input: Box<PhysicalPlan>,
        aggs: Vec<AggOp>,
        group_dims: Vec<String>,
    },
    /// Hash join on matching dimensions.
    HashJoin {
        build: Box<PhysicalPlan>,
        probe: Box<PhysicalPlan>,
        on_dims: Vec<(String, String)>,
    },
    /// Join arrays aligned on matching chunk coordinates.
    ChunkAlignedJoin {
        left: Box<PhysicalPlan>,
        right: Box<PhysicalPlan>,
        on_dims: Vec<(String, String)>,
    },
    /// Parallelism boundary — redistribute data across partitions.
    Exchange {
        input: Box<PhysicalPlan>,
        partitioning: Partitioning,
    },
    /// Rolling window aggregate along a dimension.
    RollingWindow {
        input: Box<PhysicalPlan>,
        func: AggFunc,
        attr: String,
        window_dim: usize,
        window_size: usize,
    },
    /// Block reduction by integer factors along dimensions.
    Coarsen {
        input: Box<PhysicalPlan>,
        func: AggFunc,
        factors: Vec<(usize, usize)>,
    },
    /// Concatenate two inputs along a dimension.
    Concat {
        left: Box<PhysicalPlan>,
        right: Box<PhysicalPlan>,
        along_dim: usize,
    },
    /// Compute quantile along a dimension.
    Quantile {
        input: Box<PhysicalPlan>,
        attr: String,
        quantile: f64,
        along_dim: usize,
    },
}

impl PhysicalPlan {
    /// Returns the number of child plan nodes.
    pub fn num_children(&self) -> usize {
        match self {
            Self::Scan { .. } => 0,
            Self::Trim { .. }
            | Self::Slice { .. }
            | Self::Apply { .. }
            | Self::Filter { .. }
            | Self::PartialAggregate { .. }
            | Self::FinalAggregate { .. }
            | Self::Exchange { .. }
            | Self::RollingWindow { .. }
            | Self::Coarsen { .. }
            | Self::Quantile { .. } => 1,
            Self::HashJoin { .. } | Self::ChunkAlignedJoin { .. } | Self::Concat { .. } => 2,
        }
    }
}

// ── EXPLAIN Display ──────────────────────────────────────────────────

impl fmt::Display for PhysicalPlan {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write_plan(self, f, 0)
    }
}

fn write_plan(plan: &PhysicalPlan, f: &mut fmt::Formatter<'_>, indent: usize) -> fmt::Result {
    let pad = "  ".repeat(indent);
    match plan {
        PhysicalPlan::Scan {
            array,
            source_uri,
            chunk_list,
            projection,
            use_cache,
            ..
        } => {
            writeln!(
                f,
                "{pad}Scan array={array} uri={source_uri} chunks={} proj={projection:?} cache={use_cache}",
                chunk_list.len()
            )
        }
        PhysicalPlan::Trim {
            input,
            ranges,
            global_ranges,
        } => {
            if let Some(gr) = global_ranges {
                writeln!(f, "{pad}Trim ranges={ranges:?} global={gr:?}")?;
            } else {
                writeln!(f, "{pad}Trim ranges={ranges:?}")?;
            }
            write_plan(input, f, indent + 1)
        }
        PhysicalPlan::Slice {
            input,
            axis,
            index,
            global_index,
        } => {
            if let Some(gi) = global_index {
                writeln!(f, "{pad}Slice axis={axis} index={index} global={gi}")?;
            } else {
                writeln!(f, "{pad}Slice axis={axis} index={index}")?;
            }
            write_plan(input, f, indent + 1)
        }
        PhysicalPlan::Apply {
            input,
            output_attr,
            program,
        } => {
            writeln!(
                f,
                "{pad}Apply output={output_attr} regs={} ops={}",
                program.num_registers,
                program.ops.len()
            )?;
            write_plan(input, f, indent + 1)
        }
        PhysicalPlan::Filter { input, predicate } => {
            writeln!(
                f,
                "{pad}Filter regs={} ops={}",
                predicate.num_registers,
                predicate.ops.len()
            )?;
            write_plan(input, f, indent + 1)
        }
        PhysicalPlan::PartialAggregate { input, aggs } => {
            let agg_strs: Vec<_> = aggs.iter().map(|a| a.to_string()).collect();
            writeln!(f, "{pad}PartialAggregate [{}]", agg_strs.join(", "))?;
            write_plan(input, f, indent + 1)
        }
        PhysicalPlan::FinalAggregate {
            input,
            aggs,
            group_dims,
        } => {
            let agg_strs: Vec<_> = aggs.iter().map(|a| a.to_string()).collect();
            writeln!(
                f,
                "{pad}FinalAggregate [{}] group_by={group_dims:?}",
                agg_strs.join(", ")
            )?;
            write_plan(input, f, indent + 1)
        }
        PhysicalPlan::HashJoin {
            build,
            probe,
            on_dims,
        } => {
            writeln!(f, "{pad}HashJoin on={on_dims:?}")?;
            writeln!(f, "{pad}  build:")?;
            write_plan(build, f, indent + 2)?;
            writeln!(f, "{pad}  probe:")?;
            write_plan(probe, f, indent + 2)
        }
        PhysicalPlan::ChunkAlignedJoin {
            left,
            right,
            on_dims,
        } => {
            writeln!(f, "{pad}ChunkAlignedJoin on={on_dims:?}")?;
            writeln!(f, "{pad}  left:")?;
            write_plan(left, f, indent + 2)?;
            writeln!(f, "{pad}  right:")?;
            write_plan(right, f, indent + 2)
        }
        PhysicalPlan::Exchange {
            input,
            partitioning,
        } => {
            writeln!(f, "{pad}Exchange {partitioning}")?;
            write_plan(input, f, indent + 1)
        }
        PhysicalPlan::RollingWindow {
            input,
            func,
            attr,
            window_dim,
            window_size,
        } => {
            writeln!(
                f,
                "{pad}RollingWindow func={func} attr={attr} dim={window_dim} size={window_size}"
            )?;
            write_plan(input, f, indent + 1)
        }
        PhysicalPlan::Coarsen {
            input,
            func,
            factors,
        } => {
            let factor_strs: Vec<_> = factors
                .iter()
                .map(|(axis, fac)| format!("({axis},{fac})"))
                .collect();
            writeln!(
                f,
                "{pad}Coarsen func={func} factors=[{}]",
                factor_strs.join(", ")
            )?;
            write_plan(input, f, indent + 1)
        }
        PhysicalPlan::Concat {
            left,
            right,
            along_dim,
        } => {
            writeln!(f, "{pad}Concat along_dim={along_dim}")?;
            writeln!(f, "{pad}  left:")?;
            write_plan(left, f, indent + 2)?;
            writeln!(f, "{pad}  right:")?;
            write_plan(right, f, indent + 2)
        }
        PhysicalPlan::Quantile {
            input,
            attr,
            quantile,
            along_dim,
        } => {
            writeln!(f, "{pad}Quantile attr={attr} q={quantile} dim={along_dim}")?;
            write_plan(input, f, indent + 1)
        }
    }
}

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

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

    fn simple_schema() -> ArraySchema {
        ArraySchema::new(
            "climate",
            vec![
                Dimension::new("x", DType::Float64, 0, 99, 10),
                Dimension::new("y", DType::Float64, 0, 99, 10),
            ],
            vec![Field {
                name: "temp".into(),
                dtype: DType::Float32,
            }],
        )
        .unwrap()
    }

    fn simple_program() -> ExprProgram {
        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: DType::Float64,
            num_registers: 2,
        }
    }

    #[test]
    fn plan_scan_display() {
        let plan = PhysicalPlan::Scan {
            array: "climate".into(),
            schema: simple_schema(),
            source_uri: "zarr:///data/climate.zarr".into(),
            chunk_list: vec![ChunkCoord::new(&[0, 0]), ChunkCoord::new(&[0, 1])],
            projection: vec![0],
            use_cache: true,
            stats: None,
        };
        let display = plan.to_string();
        assert!(display.contains("Scan"));
        assert!(display.contains("climate"));
        assert!(display.contains("chunks=2"));
        assert!(display.contains("cache=true"));
    }

    #[test]
    fn plan_tree_display() {
        let scan = PhysicalPlan::Scan {
            array: "climate".into(),
            schema: simple_schema(),
            source_uri: "zarr:///data/climate.zarr".into(),
            chunk_list: vec![ChunkCoord::new(&[0, 0])],
            projection: vec![0],
            use_cache: false,
            stats: None,
        };

        let apply = PhysicalPlan::Apply {
            input: Box::new(scan),
            program: simple_program(),
            output_attr: "temp_f".into(),
        };

        let agg = PhysicalPlan::FinalAggregate {
            input: Box::new(PhysicalPlan::PartialAggregate {
                input: Box::new(apply),
                aggs: vec![AggOp {
                    func: AggFunc::Avg,
                    attr: "temp_f".into(),
                    output_dtype: DType::Float64,
                }],
            }),
            aggs: vec![AggOp {
                func: AggFunc::Avg,
                attr: "temp_f".into(),
                output_dtype: DType::Float64,
            }],
            group_dims: vec![],
        };

        let display = agg.to_string();
        assert!(display.contains("FinalAggregate"));
        assert!(display.contains("PartialAggregate"));
        assert!(display.contains("Apply"));
        assert!(display.contains("Scan"));
    }

    #[test]
    fn plan_serde_roundtrip() {
        let scan = PhysicalPlan::Scan {
            array: "climate".into(),
            schema: simple_schema(),
            source_uri: "zarr:///data/climate.zarr".into(),
            chunk_list: vec![ChunkCoord::new(&[0, 0])],
            projection: vec![0],
            use_cache: true,
            stats: Some(PlanStats {
                est_chunks: 100,
                est_bytes_per_chunk: 65536,
            }),
        };

        let plan = PhysicalPlan::Apply {
            input: Box::new(scan),
            program: simple_program(),
            output_attr: "result".into(),
        };

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

    #[test]
    fn plan_num_children() {
        let scan = PhysicalPlan::Scan {
            array: "a".into(),
            schema: simple_schema(),
            source_uri: "zarr:///a".into(),
            chunk_list: vec![],
            projection: vec![],
            use_cache: false,
            stats: None,
        };
        assert_eq!(scan.num_children(), 0);

        let apply = PhysicalPlan::Apply {
            input: Box::new(scan.clone()),
            program: simple_program(),
            output_attr: "out".into(),
        };
        assert_eq!(apply.num_children(), 1);

        let join = PhysicalPlan::HashJoin {
            build: Box::new(scan.clone()),
            probe: Box::new(scan),
            on_dims: vec![("x".into(), "x".into())],
        };
        assert_eq!(join.num_children(), 2);
    }

    #[test]
    fn partitioning_display() {
        assert_eq!(Partitioning::RoundRobin(4).to_string(), "RoundRobin(4)");
        assert_eq!(Partitioning::ChunkHash(8).to_string(), "ChunkHash(8)");
        assert_eq!(
            Partitioning::DimensionRange {
                dim_name: "time".into(),
                num_partitions: 12
            }
            .to_string(),
            "DimensionRange(time, 12)"
        );
    }

    #[test]
    fn partitioning_num_partitions() {
        assert_eq!(Partitioning::RoundRobin(4).num_partitions(), 4);
        assert_eq!(Partitioning::ChunkHash(8).num_partitions(), 8);
        assert_eq!(
            Partitioning::DimensionRange {
                dim_name: "t".into(),
                num_partitions: 3
            }
            .num_partitions(),
            3
        );
    }

    #[test]
    fn aggfunc_display() {
        assert_eq!(AggFunc::Sum.to_string(), "SUM");
        assert_eq!(AggFunc::Avg.to_string(), "AVG");
        assert_eq!(AggFunc::Min.to_string(), "MIN");
        assert_eq!(AggFunc::Max.to_string(), "MAX");
        assert_eq!(AggFunc::Count.to_string(), "COUNT");
    }

    #[test]
    fn exchange_display() {
        let plan = PhysicalPlan::Exchange {
            input: Box::new(PhysicalPlan::Scan {
                array: "a".into(),
                schema: simple_schema(),
                source_uri: "zarr:///a".into(),
                chunk_list: vec![],
                projection: vec![0],
                use_cache: false,
                stats: None,
            }),
            partitioning: Partitioning::ChunkHash(4),
        };
        let display = plan.to_string();
        assert!(display.contains("Exchange ChunkHash(4)"));
        assert!(display.contains("Scan"));
    }

    #[test]
    fn join_display() {
        let scan_a = PhysicalPlan::Scan {
            array: "a".into(),
            schema: simple_schema(),
            source_uri: "zarr:///a".into(),
            chunk_list: vec![],
            projection: vec![0],
            use_cache: false,
            stats: None,
        };
        let scan_b = PhysicalPlan::Scan {
            array: "b".into(),
            schema: simple_schema(),
            source_uri: "zarr:///b".into(),
            chunk_list: vec![],
            projection: vec![0],
            use_cache: false,
            stats: None,
        };

        let plan = PhysicalPlan::HashJoin {
            build: Box::new(scan_a),
            probe: Box::new(scan_b),
            on_dims: vec![("x".into(), "x".into()), ("y".into(), "y".into())],
        };

        let display = plan.to_string();
        assert!(display.contains("HashJoin"));
        assert!(display.contains("build:"));
        assert!(display.contains("probe:"));
    }
}
