use crate::signal::bitpack::{bit_slice, sign_extend};
use crate::signal::SignalValue;
use crate::trace::{SignalId, Trace};

/// Expression tree compiled from Python-side SignalExpr.
#[derive(Clone)]
pub enum Expr {
    Signal(SignalId),
    Const(SignalValue),
    RelEval(Box<Expr>, i64),
    Not(Box<Expr>),
    And(Box<Expr>, Box<Expr>),
    Or(Box<Expr>, Box<Expr>),
    Eq(Box<Expr>, Box<Expr>),
    Neq(Box<Expr>, Box<Expr>),
    Gt(Box<Expr>, Box<Expr>),
    Lt(Box<Expr>, Box<Expr>),
    Gte(Box<Expr>, Box<Expr>),
    Lte(Box<Expr>, Box<Expr>),
    Add(Box<Expr>, Box<Expr>),
    Sub(Box<Expr>, Box<Expr>),
    BitAnd(Box<Expr>, Box<Expr>),
    BitOr(Box<Expr>, Box<Expr>),
    BitXor(Box<Expr>, Box<Expr>),
    Slice(Box<Expr>, u32, u32),
    Signed(Box<Expr>),
    True,
}

/// Evaluate an expression at a given index, returning a SignalValue.
pub fn eval(trace: &dyn Trace, expr: &Expr, index: usize) -> SignalValue {
    match expr {
        Expr::Signal(id) => {
            trace.signal_value_by_id(*id, index)
        }
        Expr::Const(v) => v.clone(),
        Expr::True => SignalValue::Scalar(1),
        Expr::RelEval(e, offset) => {
            let new_idx = index as i64 + offset;
            if new_idx < 0 || new_idx as usize > trace.max_index() {
                SignalValue::Unknown
            } else {
                eval(trace, e, new_idx as usize)
            }
        }
        Expr::Not(e) => {
            let v = eval(trace, e, index);
            SignalValue::Scalar(if v.is_truthy() { 0 } else { 1 })
        }
        Expr::And(a, b) => {
            let va = eval(trace, a, index);
            let vb = eval(trace, b, index);
            SignalValue::Scalar(if va.is_truthy() && vb.is_truthy() { 1 } else { 0 })
        }
        Expr::Or(a, b) => {
            let va = eval(trace, a, index);
            let vb = eval(trace, b, index);
            SignalValue::Scalar(if va.is_truthy() || vb.is_truthy() { 1 } else { 0 })
        }
        Expr::Eq(a, b) => {
            let va = eval(trace, a, index);
            let vb = eval(trace, b, index);
            SignalValue::Scalar(if va == vb { 1 } else {
                // Also compare numerically
                match (va.to_u64(), vb.to_u64()) {
                    (Some(x), Some(y)) if x == y => 1,
                    _ => 0,
                }
            })
        }
        Expr::Neq(a, b) => {
            let eq = eval(trace, &Expr::Eq(Box::new(Expr::Const(eval(trace, a, index))),
                                            Box::new(Expr::Const(eval(trace, b, index)))), index);
            SignalValue::Scalar(if eq.is_truthy() { 0 } else { 1 })
        }
        Expr::Gt(a, b) => binop_cmp(trace, a, b, index, |x, y| x > y),
        Expr::Lt(a, b) => binop_cmp(trace, a, b, index, |x, y| x < y),
        Expr::Gte(a, b) => binop_cmp(trace, a, b, index, |x, y| x >= y),
        Expr::Lte(a, b) => binop_cmp(trace, a, b, index, |x, y| x <= y),
        Expr::Add(a, b) => binop_arith(trace, a, b, index, |x, y| x.wrapping_add(y)),
        Expr::Sub(a, b) => binop_arith(trace, a, b, index, |x, y| x.wrapping_sub(y)),
        Expr::BitAnd(a, b) => binop_arith(trace, a, b, index, |x, y| x & y),
        Expr::BitOr(a, b) => binop_arith(trace, a, b, index, |x, y| x | y),
        Expr::BitXor(a, b) => binop_arith(trace, a, b, index, |x, y| x ^ y),
        Expr::Slice(e, high, low) => {
            let v = eval(trace, e, index);
            match v.to_u64() {
                Some(val) => SignalValue::Scalar(bit_slice(val, *high, *low)),
                None => SignalValue::Unknown,
            }
        }
        Expr::Signed(e) => {
            let v = eval(trace, e, index);
            // Find width from context - use 64 as default
            match v.to_u64() {
                Some(val) => {
                    // Try to determine width; default to 64
                    let width = 64u32;
                    let signed = sign_extend(val, width);
                    SignalValue::Scalar(signed as u64)
                }
                None => SignalValue::Unknown,
            }
        }
    }
}

fn binop_cmp(
    trace: &dyn Trace,
    a: &Expr,
    b: &Expr,
    index: usize,
    f: fn(u64, u64) -> bool,
) -> SignalValue {
    let va = eval(trace, a, index);
    let vb = eval(trace, b, index);
    match (va.to_u64(), vb.to_u64()) {
        (Some(x), Some(y)) => SignalValue::Scalar(if f(x, y) { 1 } else { 0 }),
        _ => SignalValue::Unknown,
    }
}

fn binop_arith(
    trace: &dyn Trace,
    a: &Expr,
    b: &Expr,
    index: usize,
    f: fn(u64, u64) -> u64,
) -> SignalValue {
    let va = eval(trace, a, index);
    let vb = eval(trace, b, index);
    match (va.to_u64(), vb.to_u64()) {
        (Some(x), Some(y)) => SignalValue::Scalar(f(x, y)),
        _ => SignalValue::Unknown,
    }
}

/// Evaluate an expression as a boolean at a given index.
pub fn eval_bool(trace: &dyn Trace, expr: &Expr, index: usize) -> bool {
    eval(trace, expr, index).is_truthy()
}
