use arrdb_core::{ArrDbError, ArraySchema, DType, Result, ScalarValue};
use arrdb_expr::{BinOp, ScalarExpr, UnOp, numeric_supertype, resolve_type};

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

/// Compile context — tracks register allocation and schema for name resolution.
struct CompileContext<'a> {
    schema: &'a ArraySchema,
    ops: Vec<ExprOp>,
    next_reg: u16,
}

impl<'a> CompileContext<'a> {
    fn new(schema: &'a ArraySchema) -> Self {
        Self {
            schema,
            ops: Vec::new(),
            next_reg: 0,
        }
    }

    fn alloc_reg(&mut self) -> RegId {
        let r = RegId(self.next_reg);
        self.next_reg += 1;
        r
    }

    fn emit(&mut self, op: ExprOp) {
        self.ops.push(op);
    }
}

/// Compiles a high-level `ScalarExpr` tree into a linear `ExprProgram`.
///
/// Post-order tree traversal. Fresh registers are allocated for each result
/// (no reuse in v1). Scalar optimizations are applied when one operand of a
/// binary op is a literal.
pub fn compile_expr(expr: &ScalarExpr, schema: &ArraySchema) -> Result<ExprProgram> {
    let output_dtype = arrdb_expr::resolve_type(expr, schema)?;
    let mut ctx = CompileContext::new(schema);
    let output_reg = compile_node(expr, &mut ctx)?;

    Ok(ExprProgram {
        ops: ctx.ops,
        output_reg,
        output_dtype,
        num_registers: ctx.next_reg,
    })
}

/// Recursively compile a single expression node, returning the register
/// that holds its result.
fn compile_node(expr: &ScalarExpr, ctx: &mut CompileContext<'_>) -> Result<RegId> {
    match expr {
        ScalarExpr::Literal(value) => {
            let dst = ctx.alloc_reg();
            ctx.emit(ExprOp::LoadScalar {
                dst,
                value: value.clone(),
            });
            Ok(dst)
        }

        ScalarExpr::ColumnRef(name) => {
            // Check attributes first (the common case for per-chunk expressions).
            if let Some(idx) = ctx.schema.attribute_index(name) {
                let dst = ctx.alloc_reg();
                ctx.emit(ExprOp::LoadAttr { dst, attr_idx: idx });
                return Ok(dst);
            }

            // Dimension references cannot be used in per-chunk expressions.
            if ctx.schema.dimension(name).is_some() {
                return Err(ArrDbError::Expression(format!(
                    "dimension '{name}' cannot be used in per-chunk expressions"
                )));
            }

            Err(ArrDbError::Expression(format!(
                "unknown column '{name}' in array '{}'",
                ctx.schema.name
            )))
        }

        ScalarExpr::BinaryOp { op, left, right } => compile_binary(*op, left, right, ctx),

        ScalarExpr::UnaryOp { op, operand } => {
            let src = compile_node(operand, ctx)?;
            // Float-only unary ops (transcendentals + rounding) only have
            // Float32/Float64 kernels.  If the operand is an integer type,
            // insert a cast to Float64 so the correct kernel is dispatched.
            let src = match op {
                UnOp::Sqrt
                | UnOp::Sin
                | UnOp::Cos
                | UnOp::Log
                | UnOp::Exp
                | UnOp::Ceil
                | UnOp::Floor
                | UnOp::Round => {
                    let input_dt = resolve_type(operand, ctx.schema)?;
                    if !matches!(input_dt, DType::Float32 | DType::Float64) {
                        let cast_dst = ctx.alloc_reg();
                        ctx.emit(ExprOp::Cast {
                            dst: cast_dst,
                            src,
                            to: DType::Float64,
                        });
                        cast_dst
                    } else {
                        src
                    }
                }
                _ => src,
            };
            let dst = ctx.alloc_reg();
            ctx.emit(ExprOp::UnaryOp { op: *op, dst, src });
            Ok(dst)
        }

        ScalarExpr::Cast { expr, to } => {
            let src = compile_node(expr, ctx)?;
            let dst = ctx.alloc_reg();
            ctx.emit(ExprOp::Cast {
                dst,
                src,
                to: to.clone(),
            });
            Ok(dst)
        }

        ScalarExpr::Case { when_thens, else_ } => compile_case(when_thens, else_, ctx),

        ScalarExpr::Function { name, .. } => Err(ArrDbError::NotImplemented(format!(
            "compilation of function '{name}'"
        ))),
    }
}

/// Emit a ScalarBinaryOp with type promotion.
///
/// If the buffer register dtype differs from the scalar value type,
/// insert a Cast on the buffer and coerce the scalar so both match
/// the numeric supertype.
fn emit_scalar_binop(
    op: BinOp,
    buf_expr: &ScalarExpr,
    mut buf_reg: RegId,
    scalar: &ScalarValue,
    ctx: &mut CompileContext<'_>,
) -> Result<RegId> {
    let buf_dtype = resolve_type(buf_expr, ctx.schema)?;
    let scalar_dtype = scalar.dtype();
    let rhs = if buf_dtype != scalar_dtype {
        let target = numeric_supertype(&buf_dtype, &scalar_dtype)?;
        if buf_dtype != target {
            let cast_dst = ctx.alloc_reg();
            ctx.emit(ExprOp::Cast {
                dst: cast_dst,
                src: buf_reg,
                to: target.clone(),
            });
            buf_reg = cast_dst;
        }
        coerce_scalar(scalar, &target)
    } else {
        scalar.clone()
    };
    let dst = ctx.alloc_reg();
    ctx.emit(ExprOp::ScalarBinaryOp {
        op,
        dst,
        lhs: buf_reg,
        rhs,
    });
    Ok(dst)
}

/// Compile a binary operation with scalar optimizations.
fn compile_binary(
    op: BinOp,
    left: &ScalarExpr,
    right: &ScalarExpr,
    ctx: &mut CompileContext<'_>,
) -> Result<RegId> {
    // Check if this is a comparison operation.
    if let Some(cmp) = CmpOp::from_binop(op) {
        return compile_comparison(cmp, left, right, ctx);
    }

    // Scalar optimization: one operand is a literal.
    // For commutative ops (Add, Mul), swap literal to the right.
    let is_commutative = matches!(op, BinOp::Add | BinOp::Mul | BinOp::And | BinOp::Or);

    if right.is_literal() && !left.is_literal() {
        // Right is scalar — use ScalarBinaryOp with type promotion.
        let lhs = compile_node(left, ctx)?;
        return emit_scalar_binop(op, left, lhs, right.as_literal().unwrap(), ctx);
    }

    if left.is_literal() && !right.is_literal() && is_commutative {
        // Commutative swap: literal on left → move to right.
        let lhs = compile_node(right, ctx)?;
        return emit_scalar_binop(op, right, lhs, left.as_literal().unwrap(), ctx);
    }

    if left.is_literal() && !right.is_literal() {
        // Non-commutative (Sub, Div, Mod): must keep order.
        // Load scalar into register, then BinaryOp.
        let lhs = compile_node(left, ctx)?;
        let rhs = compile_node(right, ctx)?;
        let dst = ctx.alloc_reg();
        ctx.emit(ExprOp::BinaryOp { op, dst, lhs, rhs });
        return Ok(dst);
    }

    // General case: both non-literal or both literal.
    let lhs = compile_node(left, ctx)?;
    let rhs = compile_node(right, ctx)?;
    let dst = ctx.alloc_reg();
    ctx.emit(ExprOp::BinaryOp { op, dst, lhs, rhs });
    Ok(dst)
}

/// Coerce a ScalarValue to a target DType.
fn coerce_scalar(val: &ScalarValue, target: &DType) -> ScalarValue {
    let f = match val {
        ScalarValue::Float64(v) => *v,
        ScalarValue::Float32(v) => *v as f64,
        ScalarValue::Int64(v) => *v as f64,
        ScalarValue::Int32(v) => *v as f64,
        ScalarValue::Int16(v) => *v as f64,
        ScalarValue::Int8(v) => *v as f64,
        ScalarValue::UInt64(v) => *v as f64,
        ScalarValue::UInt32(v) => *v as f64,
        ScalarValue::UInt16(v) => *v as f64,
        ScalarValue::UInt8(v) => *v as f64,
        _ => return val.clone(),
    };
    match target {
        DType::Float64 => ScalarValue::Float64(f),
        DType::Float32 => ScalarValue::Float32(f as f32),
        DType::Int64 => ScalarValue::Int64(f as i64),
        DType::Int32 => ScalarValue::Int32(f as i32),
        DType::Int16 => ScalarValue::Int16(f as i16),
        DType::Int8 => ScalarValue::Int8(f as i8),
        DType::UInt64 => ScalarValue::UInt64(f as u64),
        DType::UInt32 => ScalarValue::UInt32(f as u32),
        DType::UInt16 => ScalarValue::UInt16(f as u16),
        DType::UInt8 => ScalarValue::UInt8(f as u8),
        _ => val.clone(),
    }
}

/// Emit a CompareScalar with type promotion.
fn emit_compare_scalar(
    cmp: CmpOp,
    buf_expr: &ScalarExpr,
    mut buf_reg: RegId,
    scalar: &ScalarValue,
    ctx: &mut CompileContext<'_>,
) -> Result<RegId> {
    let buf_dtype = resolve_type(buf_expr, ctx.schema)?;
    let scalar_dtype = scalar.dtype();
    let rhs = if buf_dtype != scalar_dtype {
        let target = numeric_supertype(&buf_dtype, &scalar_dtype)?;
        if buf_dtype != target {
            let cast_dst = ctx.alloc_reg();
            ctx.emit(ExprOp::Cast {
                dst: cast_dst,
                src: buf_reg,
                to: target.clone(),
            });
            buf_reg = cast_dst;
        }
        coerce_scalar(scalar, &target)
    } else {
        scalar.clone()
    };
    let dst = ctx.alloc_reg();
    ctx.emit(ExprOp::CompareScalar {
        op: cmp,
        dst,
        lhs: buf_reg,
        rhs,
    });
    Ok(dst)
}

/// Compile a comparison operation with scalar optimization + flip.
fn compile_comparison(
    cmp: CmpOp,
    left: &ScalarExpr,
    right: &ScalarExpr,
    ctx: &mut CompileContext<'_>,
) -> Result<RegId> {
    // (expr CMP literal) → CompareScalar
    if right.is_literal() && !left.is_literal() {
        let lhs = compile_node(left, ctx)?;
        return emit_compare_scalar(cmp, left, lhs, right.as_literal().unwrap(), ctx);
    }

    // (literal CMP expr) → flip and CompareScalar
    if left.is_literal() && !right.is_literal() {
        let lhs = compile_node(right, ctx)?;
        return emit_compare_scalar(cmp.flip(), right, lhs, left.as_literal().unwrap(), ctx);
    }

    // General case: both non-literal or both literal.
    let lhs = compile_node(left, ctx)?;
    let rhs = compile_node(right, ctx)?;
    let dst = ctx.alloc_reg();
    ctx.emit(ExprOp::Compare {
        op: cmp,
        dst,
        lhs,
        rhs,
    });
    Ok(dst)
}

/// Compile a CASE expression.
///
/// Strategy: compile ELSE first, then WHEN/THEN pairs in reverse order,
/// chaining via Select ops. The last Select's result is the output.
fn compile_case(
    when_thens: &[(ScalarExpr, ScalarExpr)],
    else_: &Option<Box<ScalarExpr>>,
    ctx: &mut CompileContext<'_>,
) -> Result<RegId> {
    // Start with the ELSE value (or NULL if absent).
    let mut current = if let Some(else_expr) = else_ {
        compile_node(else_expr, ctx)?
    } else {
        let dst = ctx.alloc_reg();
        ctx.emit(ExprOp::LoadScalar {
            dst,
            value: ScalarValue::Null(DType::Float64),
        });
        dst
    };

    // Process WHEN/THEN pairs in reverse order so earlier conditions take priority.
    for (when, then) in when_thens.iter().rev() {
        let mask = compile_node(when, ctx)?;
        let then_reg = compile_node(then, ctx)?;
        let dst = ctx.alloc_reg();
        ctx.emit(ExprOp::Select {
            dst,
            mask,
            then: then_reg,
            else_: current,
        });
        current = dst;
    }

    Ok(current)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::expr_ir::ExprOp;
    use arrdb_core::{Dimension, Field};
    use arrdb_expr::UnOp;

    fn test_schema() -> ArraySchema {
        ArraySchema::new(
            "test",
            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,
                },
                Field {
                    name: "humid".into(),
                    dtype: DType::Float32,
                },
                Field {
                    name: "flag".into(),
                    dtype: DType::Bool,
                },
            ],
        )
        .unwrap()
    }

    #[test]
    fn compile_fahrenheit_conversion() {
        // temp(Float32) * 1.8(Float64) + 32(Float64)
        // → LoadAttr, Cast(Float64), ScalarBinaryOp(Mul), ScalarBinaryOp(Add)
        let expr = ScalarExpr::binary(
            BinOp::Add,
            ScalarExpr::binary(
                BinOp::Mul,
                ScalarExpr::column("temp"),
                ScalarExpr::literal(ScalarValue::Float64(1.8)),
            ),
            ScalarExpr::literal(ScalarValue::Float64(32.0)),
        );

        let program = compile_expr(&expr, &test_schema()).unwrap();
        assert_eq!(program.ops.len(), 4);
        assert_eq!(program.num_registers, 4);

        // Op 0: LoadAttr(temp)
        assert!(matches!(
            &program.ops[0],
            ExprOp::LoadAttr { attr_idx: 0, .. }
        ));
        // Op 1: Cast Float32 → Float64
        assert!(matches!(
            &program.ops[1],
            ExprOp::Cast {
                to: DType::Float64,
                ..
            }
        ));
        // Op 2: ScalarBinaryOp(Mul, %1, 1.8)
        assert!(matches!(
            &program.ops[2],
            ExprOp::ScalarBinaryOp {
                op: BinOp::Mul,
                rhs: ScalarValue::Float64(v),
                ..
            } if (*v - 1.8).abs() < f64::EPSILON
        ));
        // Op 3: ScalarBinaryOp(Add, %2, 32.0) — no cast needed (already Float64)
        assert!(matches!(
            &program.ops[3],
            ExprOp::ScalarBinaryOp {
                op: BinOp::Add,
                rhs: ScalarValue::Float64(v),
                ..
            } if (*v - 32.0).abs() < f64::EPSILON
        ));
    }

    #[test]
    fn compile_a_plus_b() {
        // temp + humid → LoadAttr, LoadAttr, BinaryOp(Add)
        let expr = ScalarExpr::binary(
            BinOp::Add,
            ScalarExpr::column("temp"),
            ScalarExpr::column("humid"),
        );

        let program = compile_expr(&expr, &test_schema()).unwrap();
        assert_eq!(program.ops.len(), 3);
        assert!(matches!(
            &program.ops[0],
            ExprOp::LoadAttr { attr_idx: 0, .. }
        ));
        assert!(matches!(
            &program.ops[1],
            ExprOp::LoadAttr { attr_idx: 1, .. }
        ));
        assert!(matches!(
            &program.ops[2],
            ExprOp::BinaryOp { op: BinOp::Add, .. }
        ));
    }

    #[test]
    fn compile_comparison_scalar() {
        // temp(Float32) > 0.0(Float64) → LoadAttr, Cast(Float64), CompareScalar(Gt)
        let expr = ScalarExpr::binary(
            BinOp::Gt,
            ScalarExpr::column("temp"),
            ScalarExpr::literal(ScalarValue::Float64(0.0)),
        );

        let program = compile_expr(&expr, &test_schema()).unwrap();
        assert_eq!(program.ops.len(), 3);
        assert!(matches!(
            &program.ops[0],
            ExprOp::LoadAttr { attr_idx: 0, .. }
        ));
        assert!(matches!(
            &program.ops[1],
            ExprOp::Cast {
                to: DType::Float64,
                ..
            }
        ));
        assert!(matches!(
            &program.ops[2],
            ExprOp::CompareScalar { op: CmpOp::Gt, .. }
        ));
        assert_eq!(program.output_dtype, DType::Bool);
    }

    #[test]
    fn compile_comparison_flip() {
        // 0.0(Float64) < temp(Float32) → flipped to: temp > 0.0, with Cast
        let expr = ScalarExpr::binary(
            BinOp::Lt,
            ScalarExpr::literal(ScalarValue::Float64(0.0)),
            ScalarExpr::column("temp"),
        );

        let program = compile_expr(&expr, &test_schema()).unwrap();
        assert_eq!(program.ops.len(), 3);
        assert!(matches!(
            &program.ops[1],
            ExprOp::Cast {
                to: DType::Float64,
                ..
            }
        ));
        assert!(matches!(
            &program.ops[2],
            ExprOp::CompareScalar {
                op: CmpOp::Gt, // Lt flipped to Gt
                ..
            }
        ));
    }

    #[test]
    fn compile_cast() {
        // CAST(temp AS Float64) → LoadAttr, Cast
        let expr = ScalarExpr::cast(ScalarExpr::column("temp"), DType::Float64);

        let program = compile_expr(&expr, &test_schema()).unwrap();
        assert_eq!(program.ops.len(), 2);
        assert!(matches!(
            &program.ops[1],
            ExprOp::Cast {
                to: DType::Float64,
                ..
            }
        ));
        assert_eq!(program.output_dtype, DType::Float64);
    }

    #[test]
    fn compile_case_expression() {
        // CASE WHEN temp > 0 THEN 1.0 ELSE 0.0 END
        let expr = ScalarExpr::Case {
            when_thens: vec![(
                ScalarExpr::binary(
                    BinOp::Gt,
                    ScalarExpr::column("temp"),
                    ScalarExpr::literal(ScalarValue::Float32(0.0)),
                ),
                ScalarExpr::literal(ScalarValue::Float64(1.0)),
            )],
            else_: Some(Box::new(ScalarExpr::literal(ScalarValue::Float64(0.0)))),
        };

        let program = compile_expr(&expr, &test_schema()).unwrap();

        // Should contain: LoadScalar(0.0), LoadAttr, CompareScalar, LoadScalar(1.0), Select
        let has_select = program
            .ops
            .iter()
            .any(|op| matches!(op, ExprOp::Select { .. }));
        assert!(has_select, "CASE should produce a Select op");
    }

    #[test]
    fn compile_nested_expression() {
        // (temp + humid) * sqrt(temp)
        let expr = ScalarExpr::binary(
            BinOp::Mul,
            ScalarExpr::binary(
                BinOp::Add,
                ScalarExpr::column("temp"),
                ScalarExpr::column("humid"),
            ),
            ScalarExpr::unary(UnOp::Sqrt, ScalarExpr::column("temp")),
        );

        let program = compile_expr(&expr, &test_schema()).unwrap();

        // Verify register wiring: the final BinaryOp(Mul) should reference
        // the outputs of (temp + humid) and sqrt(temp).
        let last = program.ops.last().unwrap();
        assert!(matches!(last, ExprOp::BinaryOp { op: BinOp::Mul, .. }));
    }

    #[test]
    fn compile_unknown_column_error() {
        let expr = ScalarExpr::column("missing");
        let err = compile_expr(&expr, &test_schema()).unwrap_err();
        assert!(err.to_string().contains("unknown column 'missing'"));
    }

    #[test]
    fn compile_dimension_ref_error() {
        let expr = ScalarExpr::column("x");
        let err = compile_expr(&expr, &test_schema()).unwrap_err();
        assert!(
            err.to_string()
                .contains("cannot be used in per-chunk expressions")
        );
    }

    #[test]
    fn compile_function_not_implemented() {
        let expr = ScalarExpr::Function {
            name: "my_func".into(),
            args: vec![ScalarExpr::column("temp")],
        };
        let err = compile_expr(&expr, &test_schema()).unwrap_err();
        assert!(matches!(err, ArrDbError::NotImplemented(_)));
    }

    #[test]
    fn compile_commutative_swap() {
        // 1.8(Float64) * temp(Float32) → swapped to: temp * 1.8, with Cast
        let expr = ScalarExpr::binary(
            BinOp::Mul,
            ScalarExpr::literal(ScalarValue::Float64(1.8)),
            ScalarExpr::column("temp"),
        );

        let program = compile_expr(&expr, &test_schema()).unwrap();
        assert_eq!(program.ops.len(), 3); // LoadAttr + Cast + ScalarBinaryOp
        assert!(matches!(
            &program.ops[1],
            ExprOp::Cast {
                to: DType::Float64,
                ..
            }
        ));
        assert!(matches!(
            &program.ops[2],
            ExprOp::ScalarBinaryOp { op: BinOp::Mul, .. }
        ));
    }

    #[test]
    fn compile_non_commutative_no_swap() {
        // 10.0 - temp → must NOT swap (subtraction is non-commutative)
        let expr = ScalarExpr::binary(
            BinOp::Sub,
            ScalarExpr::literal(ScalarValue::Float64(10.0)),
            ScalarExpr::column("temp"),
        );

        let program = compile_expr(&expr, &test_schema()).unwrap();
        // Should be: LoadScalar(10.0), LoadAttr(temp), BinaryOp(Sub)
        assert_eq!(program.ops.len(), 3);
        assert!(matches!(&program.ops[0], ExprOp::LoadScalar { .. }));
        assert!(matches!(
            &program.ops[1],
            ExprOp::LoadAttr { attr_idx: 0, .. }
        ));
        assert!(matches!(
            &program.ops[2],
            ExprOp::BinaryOp { op: BinOp::Sub, .. }
        ));
    }

    #[test]
    fn compile_output_dtype_matches_resolve_type() {
        let schema = test_schema();

        // Float32 attr → Float32
        let expr = ScalarExpr::column("temp");
        // resolve_type gives Float32 but compile_expr will error on dimension refs
        // For column refs that are attributes, it should work
        // Actually, wait - the compiler errors on dimension refs, but resolve_type doesn't
        // So for attributes it should match
        let program = compile_expr(&expr, &schema).unwrap();
        let expected = arrdb_expr::resolve_type(&expr, &schema).unwrap();
        assert_eq!(program.output_dtype, expected);

        // Comparison → Bool
        let expr = ScalarExpr::binary(
            BinOp::Gt,
            ScalarExpr::column("temp"),
            ScalarExpr::literal(ScalarValue::Float32(0.0)),
        );
        let program = compile_expr(&expr, &schema).unwrap();
        assert_eq!(program.output_dtype, DType::Bool);

        // Cast → target type
        let expr = ScalarExpr::cast(ScalarExpr::column("temp"), DType::Float64);
        let program = compile_expr(&expr, &schema).unwrap();
        assert_eq!(program.output_dtype, DType::Float64);
    }

    #[test]
    fn compile_case_multi_when() {
        // CASE WHEN temp > 100 THEN 3 WHEN temp > 50 THEN 2 ELSE 1 END
        let expr = ScalarExpr::Case {
            when_thens: vec![
                (
                    ScalarExpr::binary(
                        BinOp::Gt,
                        ScalarExpr::column("temp"),
                        ScalarExpr::literal(ScalarValue::Float32(100.0)),
                    ),
                    ScalarExpr::literal(ScalarValue::Int32(3)),
                ),
                (
                    ScalarExpr::binary(
                        BinOp::Gt,
                        ScalarExpr::column("temp"),
                        ScalarExpr::literal(ScalarValue::Float32(50.0)),
                    ),
                    ScalarExpr::literal(ScalarValue::Int32(2)),
                ),
            ],
            else_: Some(Box::new(ScalarExpr::literal(ScalarValue::Int32(1)))),
        };

        let program = compile_expr(&expr, &test_schema()).unwrap();

        // Should have 2 Select ops (one per WHEN/THEN pair).
        let select_count = program
            .ops
            .iter()
            .filter(|op| matches!(op, ExprOp::Select { .. }))
            .count();
        assert_eq!(select_count, 2);
    }

    #[test]
    fn compile_literal_only() {
        let expr = ScalarExpr::literal(ScalarValue::Int32(42));
        let program = compile_expr(&expr, &test_schema()).unwrap();
        assert_eq!(program.ops.len(), 1);
        assert_eq!(program.num_registers, 1);
        assert_eq!(program.output_dtype, DType::Int32);
    }

    #[test]
    fn compile_unary_ops() {
        // -temp
        let expr = ScalarExpr::unary(UnOp::Neg, ScalarExpr::column("temp"));
        let program = compile_expr(&expr, &test_schema()).unwrap();
        assert_eq!(program.ops.len(), 2);
        assert!(matches!(
            &program.ops[1],
            ExprOp::UnaryOp { op: UnOp::Neg, .. }
        ));

        // sqrt(temp) — Float32 input, no Cast needed
        let expr = ScalarExpr::unary(UnOp::Sqrt, ScalarExpr::column("temp"));
        let program = compile_expr(&expr, &test_schema()).unwrap();
        assert_eq!(program.ops.len(), 2); // LoadAttr + UnaryOp only
        assert!(matches!(
            &program.ops[1],
            ExprOp::UnaryOp { op: UnOp::Sqrt, .. }
        ));
    }

    #[test]
    fn compile_unary_float_cast_for_integers() {
        // Schema with an Int32 attribute to exercise the Cast insertion path.
        let schema = ArraySchema::new(
            "test",
            vec![Dimension::new("x", DType::Float64, 0, 99, 10)],
            vec![
                Field {
                    name: "temp".into(),
                    dtype: DType::Float32,
                },
                Field {
                    name: "count".into(),
                    dtype: DType::Int32,
                },
            ],
        )
        .unwrap();

        // sqrt(count) — Int32 input must get Cast to Float64 before UnaryOp.
        let expr = ScalarExpr::unary(UnOp::Sqrt, ScalarExpr::column("count"));
        let program = compile_expr(&expr, &schema).unwrap();
        // Expect: LoadAttr, Cast(Float64), UnaryOp(Sqrt)
        assert_eq!(program.ops.len(), 3);
        assert!(matches!(
            &program.ops[1],
            ExprOp::Cast {
                to: DType::Float64,
                ..
            }
        ));
        assert!(matches!(
            &program.ops[2],
            ExprOp::UnaryOp { op: UnOp::Sqrt, .. }
        ));
        assert_eq!(program.output_dtype, DType::Float64);

        // ceil(count) — same Cast insertion for rounding ops.
        let expr = ScalarExpr::unary(UnOp::Ceil, ScalarExpr::column("count"));
        let program = compile_expr(&expr, &schema).unwrap();
        assert_eq!(program.ops.len(), 3);
        assert!(matches!(
            &program.ops[1],
            ExprOp::Cast {
                to: DType::Float64,
                ..
            }
        ));
        assert!(matches!(
            &program.ops[2],
            ExprOp::UnaryOp { op: UnOp::Ceil, .. }
        ));

        // sqrt(temp) — Float32 input, no Cast.
        let expr = ScalarExpr::unary(UnOp::Sqrt, ScalarExpr::column("temp"));
        let program = compile_expr(&expr, &schema).unwrap();
        assert_eq!(program.ops.len(), 2); // LoadAttr + UnaryOp only
        assert!(matches!(
            &program.ops[1],
            ExprOp::UnaryOp { op: UnOp::Sqrt, .. }
        ));
        assert_eq!(program.output_dtype, DType::Float32);
    }
}
