use std::collections::HashMap;

use arrow::datatypes::DataType;
use serde::{Deserialize, Serialize};

use bt_kernel::BtError;
type Result<T> = std::result::Result<T, BtError>;

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub enum ScalarValue {
    Null,
    NaN,
    Bool(bool),
    Int64(i64),
    Float64(f64),
    Utf8(String),
}

impl ScalarValue {
    pub fn data_type(&self) -> DataType {
        match self {
            Self::Null => DataType::Null,
            Self::NaN => DataType::Float64,
            Self::Bool(_) => DataType::Boolean,
            Self::Int64(_) => DataType::Int64,
            Self::Float64(_) => DataType::Float64,
            Self::Utf8(_) => DataType::Utf8,
        }
    }

    pub fn as_f64(&self) -> Option<f64> {
        match self {
            Self::NaN => Some(f64::NAN),
            Self::Int64(value) => Some(*value as f64),
            Self::Float64(value) => Some(*value),
            _ => None,
        }
    }

    pub fn as_bool(&self) -> Option<bool> {
        match self {
            Self::Bool(value) => Some(*value),
            _ => None,
        }
    }

    pub fn as_utf8(&self) -> Option<&str> {
        match self {
            Self::Utf8(value) => Some(value.as_str()),
            _ => None,
        }
    }
}

impl From<bool> for ScalarValue {
    fn from(value: bool) -> Self {
        Self::Bool(value)
    }
}

impl From<i64> for ScalarValue {
    fn from(value: i64) -> Self {
        Self::Int64(value)
    }
}

impl From<f64> for ScalarValue {
    fn from(value: f64) -> Self {
        Self::Float64(value)
    }
}

impl From<String> for ScalarValue {
    fn from(value: String) -> Self {
        Self::Utf8(value)
    }
}

impl From<&str> for ScalarValue {
    fn from(value: &str) -> Self {
        Self::Utf8(value.to_string())
    }
}

/// A period/window argument that is either a fixed integer or a named parameter.
/// `#[serde(untagged)]` means JSON `14` → Fixed(14), JSON `"fast"` → Param("fast").
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum DynPeriod {
    Fixed(usize),
    Param(String),
}

impl DynPeriod {
    pub fn as_fixed(&self) -> Option<usize> {
        match self {
            Self::Fixed(v) => Some(*v),
            _ => None,
        }
    }

    pub fn is_param(&self) -> bool {
        matches!(self, Self::Param(_))
    }

    pub fn resolve(&self, params: &HashMap<String, ScalarValue>) -> Result<Self> {
        match self {
            Self::Fixed(_) => Ok(self.clone()),
            Self::Param(name) => {
                match params.get(name) {
                    None | Some(ScalarValue::Null) => Ok(Self::Fixed(2)), // sentinel for initial compile
                    Some(v) => {
                        let f = v.as_f64().ok_or_else(|| {
                            BtError::Expression(format!("period param '{name}' must be numeric"))
                        })?;
                        Ok(Self::Fixed(f as usize))
                    }
                }
            }
        }
    }
}

impl From<usize> for DynPeriod {
    fn from(v: usize) -> Self {
        Self::Fixed(v)
    }
}

/// A float argument (span, multiplier, etc.) that is either fixed or a named parameter.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum DynFloat {
    Fixed(f64),
    Param(String),
}

impl DynFloat {
    pub fn as_fixed(&self) -> Option<f64> {
        match self {
            Self::Fixed(v) => Some(*v),
            _ => None,
        }
    }

    pub fn is_param(&self) -> bool {
        matches!(self, Self::Param(_))
    }

    pub fn resolve(&self, params: &HashMap<String, ScalarValue>) -> Result<Self> {
        match self {
            Self::Fixed(_) => Ok(self.clone()),
            Self::Param(name) => {
                match params.get(name) {
                    None | Some(ScalarValue::Null) => Ok(Self::Fixed(2.0)), // sentinel for initial compile
                    Some(v) => {
                        let f = v.as_f64().ok_or_else(|| {
                            BtError::Expression(format!("float param '{name}' must be numeric"))
                        })?;
                        Ok(Self::Fixed(f))
                    }
                }
            }
        }
    }
}

impl From<f64> for DynFloat {
    fn from(v: f64) -> Self {
        Self::Fixed(v)
    }
}

#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
pub enum Expr {
    Column(String),
    Literal(ScalarValue),
    Parameter(String),

    Add(Box<Expr>, Box<Expr>),
    Sub(Box<Expr>, Box<Expr>),
    Mul(Box<Expr>, Box<Expr>),
    Div(Box<Expr>, Box<Expr>),

    Gt(Box<Expr>, Box<Expr>),
    Lt(Box<Expr>, Box<Expr>),
    Eq(Box<Expr>, Box<Expr>),
    And(Box<Expr>, Box<Expr>),
    Or(Box<Expr>, Box<Expr>),
    Not(Box<Expr>),

    Lag(Box<Expr>, DynPeriod),
    Lead(Box<Expr>, DynPeriod),
    RollingMean(Box<Expr>, DynPeriod),
    RollingStd(Box<Expr>, DynPeriod),
    RollingSum(Box<Expr>, DynPeriod),
    RollingMin(Box<Expr>, DynPeriod),
    RollingMax(Box<Expr>, DynPeriod),
    EwmMean(Box<Expr>, DynFloat),
    Diff(Box<Expr>, DynPeriod),
    PctChange(Box<Expr>, DynPeriod),
    CumSum(Box<Expr>),
    CumProd(Box<Expr>),
    Rank(Box<Expr>),
    ZScore(Box<Expr>, DynPeriod),
    Rsi(Box<Expr>, DynPeriod),
    Atr(Box<Expr>, Box<Expr>, Box<Expr>, DynPeriod),
    LinRegSlope(Box<Expr>, DynPeriod),
    LinRegValue(Box<Expr>, DynPeriod),
    LinRegR2(Box<Expr>, DynPeriod),

    // -- Moving Averages --
    Dema(Box<Expr>, DynPeriod),
    Tema(Box<Expr>, DynPeriod),
    Wma(Box<Expr>, DynPeriod),
    Hma(Box<Expr>, DynPeriod),
    Kama(Box<Expr>, DynPeriod),

    // -- Momentum --
    Roc(Box<Expr>, DynPeriod),
    Macd(Box<Expr>, DynPeriod, DynPeriod),
    MacdSignal(Box<Expr>, DynPeriod, DynPeriod, DynPeriod),
    MacdHist(Box<Expr>, DynPeriod, DynPeriod, DynPeriod),
    StochK(Box<Expr>, Box<Expr>, Box<Expr>, DynPeriod),
    WilliamsR(Box<Expr>, Box<Expr>, Box<Expr>, DynPeriod),
    Cci(Box<Expr>, Box<Expr>, Box<Expr>, DynPeriod),
    Adx(Box<Expr>, Box<Expr>, Box<Expr>, DynPeriod),

    // -- Volatility --
    TrueRange(Box<Expr>, Box<Expr>, Box<Expr>),
    Natr(Box<Expr>, Box<Expr>, Box<Expr>, DynPeriod),
    BollingerUpper(Box<Expr>, DynPeriod, DynFloat),
    BollingerLower(Box<Expr>, DynPeriod, DynFloat),
    BollingerWidth(Box<Expr>, DynPeriod, DynFloat),
    KeltnerUpper(Box<Expr>, Box<Expr>, Box<Expr>, DynPeriod, DynFloat),
    KeltnerLower(Box<Expr>, Box<Expr>, Box<Expr>, DynPeriod, DynFloat),
    SuperTrend(Box<Expr>, Box<Expr>, Box<Expr>, DynPeriod, DynFloat),

    // -- Volume --
    Obv(Box<Expr>, Box<Expr>),
    Vwap(Box<Expr>, Box<Expr>, Box<Expr>, Box<Expr>),
    AdLine(Box<Expr>, Box<Expr>, Box<Expr>, Box<Expr>),
    Mfi(Box<Expr>, Box<Expr>, Box<Expr>, Box<Expr>, DynPeriod),

    // -- Signals --
    CrossAbove(Box<Expr>, Box<Expr>),
    CrossBelow(Box<Expr>, Box<Expr>),

    // -- Statistics --
    RollingMedian(Box<Expr>, DynPeriod),

    // -- Trend --
    ParabolicSar(Box<Expr>, Box<Expr>, DynFloat, DynFloat),

    IfElse(Box<Expr>, Box<Expr>, Box<Expr>),

    CrossSectionalMean(Box<Expr>),
    CrossSectionalRank(Box<Expr>),

    /// Reference a specific symbol's column data from another symbol's context.
    /// First field is the symbol name (e.g. "BTCUSDT"), second is the column
    /// expression (must be a plain `Column` reference).
    SymbolRef(String, Box<Expr>),

    // Datetime extraction (input: timestamp nanosecond column → Float64)
    Hour(Box<Expr>),
    Minute(Box<Expr>),
    DayOfWeek(Box<Expr>),
    Month(Box<Expr>),
    DayOfMonth(Box<Expr>),

    Function(String, Vec<Expr>),

    /// Stateful scan (fold) over the time axis.
    /// Compiles to a flat register-based scalar VM executed per-row in Rust.
    Scan {
        state_names: Vec<String>,
        init_exprs: Vec<Box<Expr>>,
        update_names: Vec<String>,
        update_exprs: Vec<Box<Expr>>,
        output: String,
    },

    /// Reference a state variable's value at t-1 (only valid inside Scan update_exprs).
    ScanPrev(String),

    /// Reference a variable computed earlier in the current scan step (only inside Scan update_exprs).
    ScanVar(String),
}

// ---------------------------------------------------------------------------
// Dynamic period resolution
// ---------------------------------------------------------------------------

/// Returns true if any period/float field in the expression tree is a Param.
pub fn has_dynamic_periods(expr: &Expr) -> bool {
    macro_rules! any_dyn {
        ($($p:expr),*) => { $($p.is_param() ||)* false };
    }
    macro_rules! any_child {
        ($($e:expr),*) => { $(has_dynamic_periods($e) ||)* false };
    }
    match expr {
        // Leaf nodes
        Expr::Column(_)
        | Expr::Literal(_)
        | Expr::Parameter(_)
        | Expr::ScanPrev(_)
        | Expr::ScanVar(_)
        | Expr::Function(_, _) => false,

        // Binary ops (no periods)
        Expr::Add(a, b)
        | Expr::Sub(a, b)
        | Expr::Mul(a, b)
        | Expr::Div(a, b)
        | Expr::Gt(a, b)
        | Expr::Lt(a, b)
        | Expr::Eq(a, b)
        | Expr::And(a, b)
        | Expr::Or(a, b)
        | Expr::Obv(a, b)
        | Expr::CrossAbove(a, b)
        | Expr::CrossBelow(a, b) => any_child!(a, b),

        // Unary (no periods)
        Expr::Not(a)
        | Expr::CumSum(a)
        | Expr::CumProd(a)
        | Expr::Rank(a)
        | Expr::CrossSectionalMean(a)
        | Expr::CrossSectionalRank(a)
        | Expr::Hour(a)
        | Expr::Minute(a)
        | Expr::DayOfWeek(a)
        | Expr::Month(a)
        | Expr::DayOfMonth(a) => has_dynamic_periods(a),

        // Expr + DynPeriod
        Expr::Lag(e, p)
        | Expr::Lead(e, p)
        | Expr::RollingMean(e, p)
        | Expr::RollingStd(e, p)
        | Expr::RollingSum(e, p)
        | Expr::RollingMin(e, p)
        | Expr::RollingMax(e, p)
        | Expr::Diff(e, p)
        | Expr::PctChange(e, p)
        | Expr::ZScore(e, p)
        | Expr::Rsi(e, p)
        | Expr::LinRegSlope(e, p)
        | Expr::LinRegValue(e, p)
        | Expr::LinRegR2(e, p)
        | Expr::Dema(e, p)
        | Expr::Tema(e, p)
        | Expr::Wma(e, p)
        | Expr::Hma(e, p)
        | Expr::Kama(e, p)
        | Expr::Roc(e, p)
        | Expr::RollingMedian(e, p) => any_dyn!(p) || any_child!(e),

        // Expr + DynFloat
        Expr::EwmMean(e, f) => any_dyn!(f) || any_child!(e),

        // Expr + DynPeriod + DynPeriod
        Expr::Macd(e, p1, p2) => any_dyn!(p1, p2) || any_child!(e),

        // Expr + 3×DynPeriod
        Expr::MacdSignal(e, p1, p2, p3) | Expr::MacdHist(e, p1, p2, p3) => {
            any_dyn!(p1, p2, p3) || any_child!(e)
        }

        // HLC + DynPeriod
        Expr::Atr(h, l, c, p)
        | Expr::StochK(h, l, c, p)
        | Expr::WilliamsR(h, l, c, p)
        | Expr::Cci(h, l, c, p)
        | Expr::Adx(h, l, c, p)
        | Expr::Natr(h, l, c, p) => any_dyn!(p) || any_child!(h, l, c),

        // Expr + DynPeriod + DynFloat
        Expr::BollingerUpper(e, p, f)
        | Expr::BollingerLower(e, p, f)
        | Expr::BollingerWidth(e, p, f) => any_dyn!(p, f) || any_child!(e),

        // HLC + DynPeriod + DynFloat
        Expr::KeltnerUpper(h, l, c, p, f)
        | Expr::KeltnerLower(h, l, c, p, f)
        | Expr::SuperTrend(h, l, c, p, f) => any_dyn!(p, f) || any_child!(h, l, c),

        // HLCV + DynPeriod
        Expr::Mfi(h, l, c, v, p) => any_dyn!(p) || any_child!(h, l, c, v),

        // HLC no periods
        Expr::TrueRange(h, l, c) => any_child!(h, l, c),

        // HLCV no periods
        Expr::Vwap(a, b, c, d) | Expr::AdLine(a, b, c, d) => any_child!(a, b, c, d),

        // Expr + 2×DynFloat
        Expr::ParabolicSar(h, l, f1, f2) => any_dyn!(f1, f2) || any_child!(h, l),

        // Special
        Expr::IfElse(cond, t, f) => any_child!(cond, t, f),
        Expr::SymbolRef(_, e) => has_dynamic_periods(e),
        Expr::Scan {
            init_exprs,
            update_exprs,
            ..
        } => {
            init_exprs.iter().any(|e| has_dynamic_periods(e))
                || update_exprs.iter().any(|e| has_dynamic_periods(e))
        }
    }
}

/// Resolve all `DynPeriod::Param` and `DynFloat::Param` in an expression tree.
/// Returns a new tree with all periods resolved to fixed values.
pub fn resolve_dynamic_periods(expr: &Expr, params: &HashMap<String, ScalarValue>) -> Result<Expr> {
    macro_rules! rp {
        ($p:expr) => {
            $p.resolve(params)?
        };
    }
    macro_rules! re {
        ($e:expr) => {
            Box::new(resolve_dynamic_periods($e, params)?)
        };
    }
    Ok(match expr {
        // Leaf nodes — clone as-is
        Expr::Column(s) => Expr::Column(s.clone()),
        Expr::Literal(v) => Expr::Literal(v.clone()),
        Expr::Parameter(s) => Expr::Parameter(s.clone()),
        Expr::ScanPrev(s) => Expr::ScanPrev(s.clone()),
        Expr::ScanVar(s) => Expr::ScanVar(s.clone()),
        Expr::Function(name, args) => Expr::Function(
            name.clone(),
            args.iter()
                .map(|a| resolve_dynamic_periods(a, params))
                .collect::<Result<_>>()?,
        ),

        // Binary (no periods)
        Expr::Add(a, b) => Expr::Add(re!(a), re!(b)),
        Expr::Sub(a, b) => Expr::Sub(re!(a), re!(b)),
        Expr::Mul(a, b) => Expr::Mul(re!(a), re!(b)),
        Expr::Div(a, b) => Expr::Div(re!(a), re!(b)),
        Expr::Gt(a, b) => Expr::Gt(re!(a), re!(b)),
        Expr::Lt(a, b) => Expr::Lt(re!(a), re!(b)),
        Expr::Eq(a, b) => Expr::Eq(re!(a), re!(b)),
        Expr::And(a, b) => Expr::And(re!(a), re!(b)),
        Expr::Or(a, b) => Expr::Or(re!(a), re!(b)),
        Expr::Obv(a, b) => Expr::Obv(re!(a), re!(b)),
        Expr::CrossAbove(a, b) => Expr::CrossAbove(re!(a), re!(b)),
        Expr::CrossBelow(a, b) => Expr::CrossBelow(re!(a), re!(b)),

        // Unary (no periods)
        Expr::Not(a) => Expr::Not(re!(a)),
        Expr::CumSum(a) => Expr::CumSum(re!(a)),
        Expr::CumProd(a) => Expr::CumProd(re!(a)),
        Expr::Rank(a) => Expr::Rank(re!(a)),
        Expr::CrossSectionalMean(a) => Expr::CrossSectionalMean(re!(a)),
        Expr::CrossSectionalRank(a) => Expr::CrossSectionalRank(re!(a)),
        Expr::Hour(a) => Expr::Hour(re!(a)),
        Expr::Minute(a) => Expr::Minute(re!(a)),
        Expr::DayOfWeek(a) => Expr::DayOfWeek(re!(a)),
        Expr::Month(a) => Expr::Month(re!(a)),
        Expr::DayOfMonth(a) => Expr::DayOfMonth(re!(a)),

        // Expr + DynPeriod
        Expr::Lag(e, p) => Expr::Lag(re!(e), rp!(p)),
        Expr::Lead(e, p) => Expr::Lead(re!(e), rp!(p)),
        Expr::RollingMean(e, p) => Expr::RollingMean(re!(e), rp!(p)),
        Expr::RollingStd(e, p) => Expr::RollingStd(re!(e), rp!(p)),
        Expr::RollingSum(e, p) => Expr::RollingSum(re!(e), rp!(p)),
        Expr::RollingMin(e, p) => Expr::RollingMin(re!(e), rp!(p)),
        Expr::RollingMax(e, p) => Expr::RollingMax(re!(e), rp!(p)),
        Expr::Diff(e, p) => Expr::Diff(re!(e), rp!(p)),
        Expr::PctChange(e, p) => Expr::PctChange(re!(e), rp!(p)),
        Expr::ZScore(e, p) => Expr::ZScore(re!(e), rp!(p)),
        Expr::Rsi(e, p) => Expr::Rsi(re!(e), rp!(p)),
        Expr::LinRegSlope(e, p) => Expr::LinRegSlope(re!(e), rp!(p)),
        Expr::LinRegValue(e, p) => Expr::LinRegValue(re!(e), rp!(p)),
        Expr::LinRegR2(e, p) => Expr::LinRegR2(re!(e), rp!(p)),
        Expr::Dema(e, p) => Expr::Dema(re!(e), rp!(p)),
        Expr::Tema(e, p) => Expr::Tema(re!(e), rp!(p)),
        Expr::Wma(e, p) => Expr::Wma(re!(e), rp!(p)),
        Expr::Hma(e, p) => Expr::Hma(re!(e), rp!(p)),
        Expr::Kama(e, p) => Expr::Kama(re!(e), rp!(p)),
        Expr::Roc(e, p) => Expr::Roc(re!(e), rp!(p)),
        Expr::RollingMedian(e, p) => Expr::RollingMedian(re!(e), rp!(p)),

        // Expr + DynFloat
        Expr::EwmMean(e, f) => Expr::EwmMean(re!(e), rp!(f)),

        // Expr + 2×DynPeriod
        Expr::Macd(e, p1, p2) => Expr::Macd(re!(e), rp!(p1), rp!(p2)),

        // Expr + 3×DynPeriod
        Expr::MacdSignal(e, p1, p2, p3) => Expr::MacdSignal(re!(e), rp!(p1), rp!(p2), rp!(p3)),
        Expr::MacdHist(e, p1, p2, p3) => Expr::MacdHist(re!(e), rp!(p1), rp!(p2), rp!(p3)),

        // HLC + DynPeriod
        Expr::Atr(h, l, c, p) => Expr::Atr(re!(h), re!(l), re!(c), rp!(p)),
        Expr::StochK(h, l, c, p) => Expr::StochK(re!(h), re!(l), re!(c), rp!(p)),
        Expr::WilliamsR(h, l, c, p) => Expr::WilliamsR(re!(h), re!(l), re!(c), rp!(p)),
        Expr::Cci(h, l, c, p) => Expr::Cci(re!(h), re!(l), re!(c), rp!(p)),
        Expr::Adx(h, l, c, p) => Expr::Adx(re!(h), re!(l), re!(c), rp!(p)),
        Expr::Natr(h, l, c, p) => Expr::Natr(re!(h), re!(l), re!(c), rp!(p)),

        // Expr + DynPeriod + DynFloat
        Expr::BollingerUpper(e, p, f) => Expr::BollingerUpper(re!(e), rp!(p), rp!(f)),
        Expr::BollingerLower(e, p, f) => Expr::BollingerLower(re!(e), rp!(p), rp!(f)),
        Expr::BollingerWidth(e, p, f) => Expr::BollingerWidth(re!(e), rp!(p), rp!(f)),

        // HLC + DynPeriod + DynFloat
        Expr::KeltnerUpper(h, l, c, p, f) => {
            Expr::KeltnerUpper(re!(h), re!(l), re!(c), rp!(p), rp!(f))
        }
        Expr::KeltnerLower(h, l, c, p, f) => {
            Expr::KeltnerLower(re!(h), re!(l), re!(c), rp!(p), rp!(f))
        }
        Expr::SuperTrend(h, l, c, p, f) => Expr::SuperTrend(re!(h), re!(l), re!(c), rp!(p), rp!(f)),

        // HLCV + DynPeriod
        Expr::Mfi(h, l, c, v, p) => Expr::Mfi(re!(h), re!(l), re!(c), re!(v), rp!(p)),

        // HLC no periods
        Expr::TrueRange(h, l, c) => Expr::TrueRange(re!(h), re!(l), re!(c)),

        // HLCV no periods
        Expr::Vwap(a, b, c, d) => Expr::Vwap(re!(a), re!(b), re!(c), re!(d)),
        Expr::AdLine(a, b, c, d) => Expr::AdLine(re!(a), re!(b), re!(c), re!(d)),

        // HL + 2×DynFloat
        Expr::ParabolicSar(h, l, f1, f2) => Expr::ParabolicSar(re!(h), re!(l), rp!(f1), rp!(f2)),

        // Special
        Expr::IfElse(cond, t, f) => Expr::IfElse(re!(cond), re!(t), re!(f)),
        Expr::SymbolRef(name, e) => Expr::SymbolRef(name.clone(), re!(e)),
        Expr::Scan {
            state_names,
            init_exprs,
            update_names,
            update_exprs,
            output,
        } => Expr::Scan {
            state_names: state_names.clone(),
            init_exprs: init_exprs
                .iter()
                .map(|e| Ok(Box::new(resolve_dynamic_periods(e, params)?)))
                .collect::<Result<_>>()?,
            update_names: update_names.clone(),
            update_exprs: update_exprs
                .iter()
                .map(|e| Ok(Box::new(resolve_dynamic_periods(e, params)?)))
                .collect::<Result<_>>()?,
            output: output.clone(),
        },
    })
}
