use std::fmt;

use serde::{Deserialize, Serialize};

use arrdb_core::{DType, ScalarValue};
use arrdb_expr::BinOp;

/// A virtual register — indexes into a pre-allocated buffer array.
///
/// During execution, each register maps to an `AlignedBuffer` (chunk-sized)
/// or a scalar broadcast buffer.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct RegId(pub u16);

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

/// Comparison operators — the subset of `BinOp` that produces bitmasks.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum CmpOp {
    Eq,
    Ne,
    Lt,
    Le,
    Gt,
    Ge,
}

impl CmpOp {
    /// Converts a `BinOp` comparison into a `CmpOp`.
    ///
    /// Returns `None` if the `BinOp` is not a comparison operator.
    pub fn from_binop(op: BinOp) -> Option<Self> {
        match op {
            BinOp::Eq => Some(Self::Eq),
            BinOp::Ne => Some(Self::Ne),
            BinOp::Lt => Some(Self::Lt),
            BinOp::Le => Some(Self::Le),
            BinOp::Gt => Some(Self::Gt),
            BinOp::Ge => Some(Self::Ge),
            _ => None,
        }
    }

    /// Flips the comparison for operand swap: `(a < b)` becomes `(b > a)`.
    pub fn flip(self) -> Self {
        match self {
            Self::Eq => Self::Eq,
            Self::Ne => Self::Ne,
            Self::Lt => Self::Gt,
            Self::Le => Self::Ge,
            Self::Gt => Self::Lt,
            Self::Ge => Self::Le,
        }
    }
}

impl fmt::Display for CmpOp {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Eq => write!(f, "Eq"),
            Self::Ne => write!(f, "Ne"),
            Self::Lt => write!(f, "Lt"),
            Self::Le => write!(f, "Le"),
            Self::Gt => write!(f, "Gt"),
            Self::Ge => write!(f, "Ge"),
        }
    }
}

/// Reduction operators for partial aggregation over a buffer.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum RedOp {
    Sum,
    Min,
    Max,
    Count,
    Avg,
}

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

/// A single operation in the expression program.
///
/// Each variant maps 1:1 to a SIMD kernel dispatch in `arrdb-compute`.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum ExprOp {
    /// Load an attribute buffer from the current chunk.
    LoadAttr { dst: RegId, attr_idx: usize },
    /// Load a scalar constant (broadcast to buffer size at execution time).
    LoadScalar { dst: RegId, value: ScalarValue },

    /// Element-wise binary operation on two register buffers.
    BinaryOp {
        op: BinOp,
        dst: RegId,
        lhs: RegId,
        rhs: RegId,
    },
    /// Element-wise unary operation on a register buffer.
    UnaryOp {
        op: arrdb_expr::UnOp,
        dst: RegId,
        src: RegId,
    },
    /// Element-wise binary operation: register op scalar (avoids broadcast).
    ScalarBinaryOp {
        op: BinOp,
        dst: RegId,
        lhs: RegId,
        rhs: ScalarValue,
    },

    /// Element-wise comparison of two register buffers → bitmask.
    Compare {
        op: CmpOp,
        dst: RegId,
        lhs: RegId,
        rhs: RegId,
    },
    /// Element-wise comparison: register vs scalar → bitmask.
    CompareScalar {
        op: CmpOp,
        dst: RegId,
        lhs: RegId,
        rhs: ScalarValue,
    },

    /// Cast a register buffer to a different type.
    Cast { dst: RegId, src: RegId, to: DType },
    /// Conditional select: `dst[i] = mask[i] ? then[i] : else_[i]`.
    Select {
        dst: RegId,
        mask: RegId,
        then: RegId,
        else_: RegId,
    },

    /// Reduce a buffer to a scalar (partial aggregation).
    Reduce { op: RedOp, dst: RegId, src: RegId },
}

/// A compiled expression program — linear sequence of ops + metadata.
///
/// Produced by `compile_expr()` from a `ScalarExpr` tree, consumed by the
/// SIMD kernel dispatcher in `arrdb-compute`.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct ExprProgram {
    /// The instruction sequence, executed in order.
    pub ops: Vec<ExprOp>,
    /// Which register holds the final result after execution.
    pub output_reg: RegId,
    /// The data type of the output register.
    pub output_dtype: DType,
    /// Total number of registers needed (for pre-allocation).
    pub num_registers: u16,
}

impl fmt::Display for ExprProgram {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for op in &self.ops {
            match op {
                ExprOp::LoadAttr { dst, attr_idx } => {
                    writeln!(f, "{dst} = LoadAttr(attr_idx={attr_idx})")?;
                }
                ExprOp::LoadScalar { dst, value } => {
                    writeln!(f, "{dst} = LoadScalar({value})")?;
                }
                ExprOp::BinaryOp { op, dst, lhs, rhs } => {
                    writeln!(f, "{dst} = BinaryOp({op}, {lhs}, {rhs})")?;
                }
                ExprOp::UnaryOp { op, dst, src } => {
                    writeln!(f, "{dst} = UnaryOp({op}, {src})")?;
                }
                ExprOp::ScalarBinaryOp { op, dst, lhs, rhs } => {
                    writeln!(f, "{dst} = ScalarBinaryOp({op}, {lhs}, {rhs})")?;
                }
                ExprOp::Compare { op, dst, lhs, rhs } => {
                    writeln!(f, "{dst} = Compare({op}, {lhs}, {rhs})")?;
                }
                ExprOp::CompareScalar { op, dst, lhs, rhs } => {
                    writeln!(f, "{dst} = CompareScalar({op}, {lhs}, {rhs})")?;
                }
                ExprOp::Cast { dst, src, to } => {
                    writeln!(f, "{dst} = Cast({src}, {to})")?;
                }
                ExprOp::Select {
                    dst,
                    mask,
                    then,
                    else_,
                } => {
                    writeln!(f, "{dst} = Select({mask}, {then}, {else_})")?;
                }
                ExprOp::Reduce { op, dst, src } => {
                    writeln!(f, "{dst} = Reduce({op}, {src})")?;
                }
            }
        }
        write!(f, "output = {}", self.output_reg)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn regid_display() {
        assert_eq!(RegId(0).to_string(), "%0");
        assert_eq!(RegId(42).to_string(), "%42");
    }

    #[test]
    fn cmpop_from_binop() {
        assert_eq!(CmpOp::from_binop(BinOp::Eq), Some(CmpOp::Eq));
        assert_eq!(CmpOp::from_binop(BinOp::Lt), Some(CmpOp::Lt));
        assert_eq!(CmpOp::from_binop(BinOp::Ge), Some(CmpOp::Ge));
        assert_eq!(CmpOp::from_binop(BinOp::Add), None);
        assert_eq!(CmpOp::from_binop(BinOp::And), None);
    }

    #[test]
    fn cmpop_flip() {
        assert_eq!(CmpOp::Eq.flip(), CmpOp::Eq);
        assert_eq!(CmpOp::Ne.flip(), CmpOp::Ne);
        assert_eq!(CmpOp::Lt.flip(), CmpOp::Gt);
        assert_eq!(CmpOp::Le.flip(), CmpOp::Ge);
        assert_eq!(CmpOp::Gt.flip(), CmpOp::Lt);
        assert_eq!(CmpOp::Ge.flip(), CmpOp::Le);
    }

    #[test]
    fn cmpop_flip_involution() {
        for op in [
            CmpOp::Eq,
            CmpOp::Ne,
            CmpOp::Lt,
            CmpOp::Le,
            CmpOp::Gt,
            CmpOp::Ge,
        ] {
            assert_eq!(op.flip().flip(), op);
        }
    }

    #[test]
    fn expr_program_display_snapshot() {
        // temp * 1.8 + 32
        let 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),
                },
                ExprOp::ScalarBinaryOp {
                    op: BinOp::Add,
                    dst: RegId(2),
                    lhs: RegId(1),
                    rhs: ScalarValue::Float64(32.0),
                },
            ],
            output_reg: RegId(2),
            output_dtype: DType::Float64,
            num_registers: 3,
        };

        let display = program.to_string();
        let expected = "\
%0 = LoadAttr(attr_idx=0)
%1 = ScalarBinaryOp(*, %0, 1.8)
%2 = ScalarBinaryOp(+, %1, 32)
output = %2";
        assert_eq!(display, expected);
    }

    #[test]
    fn expr_program_serde_roundtrip() {
        let program = ExprProgram {
            ops: vec![
                ExprOp::LoadAttr {
                    dst: RegId(0),
                    attr_idx: 0,
                },
                ExprOp::LoadScalar {
                    dst: RegId(1),
                    value: ScalarValue::Float64(1.8),
                },
                ExprOp::BinaryOp {
                    op: BinOp::Mul,
                    dst: RegId(2),
                    lhs: RegId(0),
                    rhs: RegId(1),
                },
                ExprOp::Compare {
                    op: CmpOp::Gt,
                    dst: RegId(3),
                    lhs: RegId(2),
                    rhs: RegId(0),
                },
                ExprOp::CompareScalar {
                    op: CmpOp::Le,
                    dst: RegId(4),
                    lhs: RegId(0),
                    rhs: ScalarValue::Float64(0.0),
                },
                ExprOp::Cast {
                    dst: RegId(5),
                    src: RegId(0),
                    to: DType::Float64,
                },
                ExprOp::Select {
                    dst: RegId(6),
                    mask: RegId(3),
                    then: RegId(2),
                    else_: RegId(5),
                },
                ExprOp::Reduce {
                    op: RedOp::Sum,
                    dst: RegId(7),
                    src: RegId(6),
                },
            ],
            output_reg: RegId(7),
            output_dtype: DType::Float64,
            num_registers: 8,
        };

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

    #[test]
    fn manual_expr_program_construction() {
        // a + b
        let program = ExprProgram {
            ops: vec![
                ExprOp::LoadAttr {
                    dst: RegId(0),
                    attr_idx: 0,
                },
                ExprOp::LoadAttr {
                    dst: RegId(1),
                    attr_idx: 1,
                },
                ExprOp::BinaryOp {
                    op: BinOp::Add,
                    dst: RegId(2),
                    lhs: RegId(0),
                    rhs: RegId(1),
                },
            ],
            output_reg: RegId(2),
            output_dtype: DType::Float32,
            num_registers: 3,
        };

        assert_eq!(program.ops.len(), 3);
        assert_eq!(program.output_reg, RegId(2));
        assert_eq!(program.num_registers, 3);
    }

    #[test]
    fn display_all_op_variants() {
        let program = ExprProgram {
            ops: vec![
                ExprOp::LoadAttr {
                    dst: RegId(0),
                    attr_idx: 0,
                },
                ExprOp::LoadScalar {
                    dst: RegId(1),
                    value: ScalarValue::Float64(0.0),
                },
                ExprOp::BinaryOp {
                    op: BinOp::Add,
                    dst: RegId(2),
                    lhs: RegId(0),
                    rhs: RegId(1),
                },
                ExprOp::UnaryOp {
                    op: arrdb_expr::UnOp::Sqrt,
                    dst: RegId(3),
                    src: RegId(2),
                },
                ExprOp::ScalarBinaryOp {
                    op: BinOp::Mul,
                    dst: RegId(4),
                    lhs: RegId(3),
                    rhs: ScalarValue::Float64(2.0),
                },
                ExprOp::Compare {
                    op: CmpOp::Gt,
                    dst: RegId(5),
                    lhs: RegId(4),
                    rhs: RegId(1),
                },
                ExprOp::CompareScalar {
                    op: CmpOp::Lt,
                    dst: RegId(6),
                    lhs: RegId(0),
                    rhs: ScalarValue::Float64(0.0),
                },
                ExprOp::Cast {
                    dst: RegId(7),
                    src: RegId(0),
                    to: DType::Int32,
                },
                ExprOp::Select {
                    dst: RegId(8),
                    mask: RegId(5),
                    then: RegId(4),
                    else_: RegId(7),
                },
                ExprOp::Reduce {
                    op: RedOp::Max,
                    dst: RegId(9),
                    src: RegId(8),
                },
            ],
            output_reg: RegId(9),
            output_dtype: DType::Float64,
            num_registers: 10,
        };

        let display = program.to_string();
        assert!(display.contains("LoadAttr(attr_idx=0)"));
        assert!(display.contains("LoadScalar(0)"));
        assert!(display.contains("BinaryOp(+, %0, %1)"));
        assert!(display.contains("UnaryOp(sqrt, %2)"));
        assert!(display.contains("ScalarBinaryOp(*, %3, 2)"));
        assert!(display.contains("Compare(Gt, %4, %1)"));
        assert!(display.contains("CompareScalar(Lt, %0, 0)"));
        assert!(display.contains("Cast(%0, INT32)"));
        assert!(display.contains("Select(%5, %4, %7)"));
        assert!(display.contains("Reduce(Max, %8)"));
        assert!(display.contains("output = %9"));
    }
}
