use std::collections::HashMap;

use bt_kernel::{BtError, Result};

use crate::expr::{Expr, ScalarValue};

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum ValueType {
    Unknown,
    Bool,
    Int64,
    Float64,
    Utf8,
}

pub fn infer_expr_type(
    expr: &Expr,
    columns: &HashMap<String, ValueType>,
    params: &HashMap<String, ScalarValue>,
) -> Result<ValueType> {
    infer_type_inner(expr, columns, params, false)
}

fn infer_type_inner(
    expr: &Expr,
    columns: &HashMap<String, ValueType>,
    params: &HashMap<String, ScalarValue>,
    in_scan: bool,
) -> Result<ValueType> {
    let recurse = |e: &Expr| infer_type_inner(e, columns, params, in_scan);

    match expr {
        Expr::Column(name) => Ok(*columns.get(name).unwrap_or(&ValueType::Unknown)),
        Expr::Literal(value) => Ok(type_of_scalar(value)),
        Expr::Parameter(name) => Ok(params
            .get(name)
            .map(type_of_scalar)
            .unwrap_or(ValueType::Unknown)),

        Expr::Add(lhs, rhs) | Expr::Sub(lhs, rhs) | Expr::Mul(lhs, rhs) | Expr::Div(lhs, rhs) => {
            let lhs = recurse(lhs)?;
            let rhs = recurse(rhs)?;
            ensure_numeric(lhs, "arithmetic lhs")?;
            ensure_numeric(rhs, "arithmetic rhs")?;
            Ok(promote_numeric(lhs, rhs))
        }

        Expr::Gt(lhs, rhs) | Expr::Lt(lhs, rhs) | Expr::Eq(lhs, rhs) => {
            let lhs = recurse(lhs)?;
            let rhs = recurse(rhs)?;
            if !is_compatible(lhs, rhs) {
                return Err(BtError::Expression(format!(
                    "type mismatch in comparison: {lhs:?} vs {rhs:?}"
                )));
            }

            Ok(ValueType::Bool)
        }

        Expr::And(lhs, rhs) | Expr::Or(lhs, rhs) => {
            let lhs = recurse(lhs)?;
            let rhs = recurse(rhs)?;
            ensure_bool(lhs, "boolean lhs")?;
            ensure_bool(rhs, "boolean rhs")?;
            Ok(ValueType::Bool)
        }
        Expr::Not(inner) => {
            let inner = recurse(inner)?;
            ensure_bool(inner, "boolean expression")?;
            Ok(ValueType::Bool)
        }

        Expr::Lag(inner, _)
        | Expr::Lead(inner, _)
        | Expr::RollingMean(inner, _)
        | Expr::RollingStd(inner, _)
        | Expr::RollingSum(inner, _)
        | Expr::RollingMin(inner, _)
        | Expr::RollingMax(inner, _)
        | Expr::EwmMean(inner, _)
        | Expr::Diff(inner, _)
        | Expr::PctChange(inner, _)
        | Expr::CumSum(inner)
        | Expr::CumProd(inner)
        | Expr::Rank(inner)
        | Expr::ZScore(inner, _)
        | Expr::Rsi(inner, _)
        | Expr::LinRegSlope(inner, _)
        | Expr::LinRegValue(inner, _)
        | Expr::LinRegR2(inner, _)
        | Expr::Dema(inner, _)
        | Expr::Tema(inner, _)
        | Expr::Wma(inner, _)
        | Expr::Hma(inner, _)
        | Expr::Kama(inner, _)
        | Expr::Roc(inner, _)
        | Expr::Macd(inner, _, _)
        | Expr::MacdSignal(inner, _, _, _)
        | Expr::MacdHist(inner, _, _, _)
        | Expr::BollingerUpper(inner, _, _)
        | Expr::BollingerLower(inner, _, _)
        | Expr::BollingerWidth(inner, _, _)
        | Expr::RollingMedian(inner, _)
        | Expr::CrossSectionalMean(inner)
        | Expr::CrossSectionalRank(inner) => {
            let inner = recurse(inner)?;
            ensure_numeric(inner, "timeseries expression")?;
            Ok(ValueType::Float64)
        }

        Expr::SymbolRef(_symbol, inner) => {
            let inner = recurse(inner)?;
            ensure_numeric(inner, "SymbolRef")?;
            Ok(ValueType::Float64)
        }

        Expr::Hour(inner)
        | Expr::Minute(inner)
        | Expr::DayOfWeek(inner)
        | Expr::Month(inner)
        | Expr::DayOfMonth(inner) => {
            let _inner = recurse(inner)?;
            Ok(ValueType::Float64)
        }

        Expr::Atr(high, low, close, _)
        | Expr::TrueRange(high, low, close)
        | Expr::Natr(high, low, close, _)
        | Expr::StochK(high, low, close, _)
        | Expr::WilliamsR(high, low, close, _)
        | Expr::Cci(high, low, close, _)
        | Expr::Adx(high, low, close, _)
        | Expr::SuperTrend(high, low, close, _, _)
        | Expr::KeltnerUpper(high, low, close, _, _)
        | Expr::KeltnerLower(high, low, close, _, _) => {
            let h = recurse(high)?;
            let l = recurse(low)?;
            let c = recurse(close)?;
            ensure_numeric(h, "indicator high")?;
            ensure_numeric(l, "indicator low")?;
            ensure_numeric(c, "indicator close")?;
            Ok(ValueType::Float64)
        }

        Expr::Vwap(high, low, close, volume)
        | Expr::AdLine(high, low, close, volume)
        | Expr::Mfi(high, low, close, volume, _) => {
            let h = recurse(high)?;
            let l = recurse(low)?;
            let c = recurse(close)?;
            let v = recurse(volume)?;
            ensure_numeric(h, "indicator high")?;
            ensure_numeric(l, "indicator low")?;
            ensure_numeric(c, "indicator close")?;
            ensure_numeric(v, "indicator volume")?;
            Ok(ValueType::Float64)
        }

        Expr::Obv(close, volume) => {
            let c = recurse(close)?;
            let v = recurse(volume)?;
            ensure_numeric(c, "obv close")?;
            ensure_numeric(v, "obv volume")?;
            Ok(ValueType::Float64)
        }

        Expr::CrossAbove(a, b) | Expr::CrossBelow(a, b) => {
            let at = recurse(a)?;
            let bt = recurse(b)?;
            ensure_numeric(at, "crossover lhs")?;
            ensure_numeric(bt, "crossover rhs")?;
            Ok(ValueType::Bool)
        }

        Expr::ParabolicSar(high, low, _, _) => {
            let h = recurse(high)?;
            let l = recurse(low)?;
            ensure_numeric(h, "parabolic_sar high")?;
            ensure_numeric(l, "parabolic_sar low")?;
            Ok(ValueType::Float64)
        }

        Expr::IfElse(condition, then_value, else_value) => {
            let condition = recurse(condition)?;
            ensure_bool(condition, "if/else condition")?;

            let then_value = recurse(then_value)?;
            let else_value = recurse(else_value)?;
            if !is_compatible(then_value, else_value) {
                return Err(BtError::Expression(format!(
                    "if/else branch type mismatch: {then_value:?} vs {else_value:?}"
                )));
            }

            Ok(promote_numeric(then_value, else_value))
        }

        Expr::Function(name, args) => {
            for arg in args {
                let ty = recurse(arg)?;
                ensure_numeric(ty, &format!("function '{name}' argument"))?;
            }

            Ok(ValueType::Float64)
        }

        Expr::Scan {
            init_exprs,
            update_exprs,
            ..
        } => {
            for (i, ie) in init_exprs.iter().enumerate() {
                let ty = recurse(ie)?;
                ensure_numeric(ty, &format!("scan init_expr[{i}]"))?;
            }
            for (i, ue) in update_exprs.iter().enumerate() {
                let ty = infer_type_inner(ue, columns, params, true)?;
                ensure_numeric(ty, &format!("scan update_expr[{i}]"))?;
            }
            Ok(ValueType::Float64)
        }

        Expr::ScanPrev(_) | Expr::ScanVar(_) => {
            if in_scan {
                Ok(ValueType::Float64)
            } else {
                Err(BtError::Expression(
                    "ScanPrev/ScanVar can only appear inside a Scan update expression".to_string(),
                ))
            }
        }
    }
}

fn type_of_scalar(value: &ScalarValue) -> ValueType {
    match value {
        ScalarValue::Null => ValueType::Unknown,
        ScalarValue::NaN => ValueType::Float64,
        ScalarValue::Bool(_) => ValueType::Bool,
        ScalarValue::Int64(_) => ValueType::Int64,
        ScalarValue::Float64(_) => ValueType::Float64,
        ScalarValue::Utf8(_) => ValueType::Utf8,
    }
}

fn ensure_numeric(value: ValueType, context: &str) -> Result<()> {
    if matches!(
        value,
        ValueType::Unknown | ValueType::Int64 | ValueType::Float64
    ) {
        return Ok(());
    }

    Err(BtError::Expression(format!(
        "{context} must be numeric, got {value:?}"
    )))
}

fn ensure_bool(value: ValueType, context: &str) -> Result<()> {
    if matches!(value, ValueType::Unknown | ValueType::Bool) {
        return Ok(());
    }

    Err(BtError::Expression(format!(
        "{context} must be bool, got {value:?}"
    )))
}

fn promote_numeric(lhs: ValueType, rhs: ValueType) -> ValueType {
    match (lhs, rhs) {
        (ValueType::Float64, _) | (_, ValueType::Float64) => ValueType::Float64,
        (ValueType::Int64, ValueType::Int64) => ValueType::Int64,
        (ValueType::Bool, ValueType::Bool) => ValueType::Bool,
        (ValueType::Utf8, ValueType::Utf8) => ValueType::Utf8,
        _ => ValueType::Float64,
    }
}

fn is_compatible(lhs: ValueType, rhs: ValueType) -> bool {
    lhs == rhs || lhs == ValueType::Unknown || rhs == ValueType::Unknown
}
