//! Unified const expression evaluation: substitution, folding, and const fn calls in one pass.

use std::collections::HashMap;

use ty_ree::types::ConstExpression;

use crate::const_eval::ConstEvaluator;

/// Context for const expression evaluation.
pub struct EvalContext<'a> {
    /// Named constant bindings (from monomorphization or const declarations).
    pub bindings: Option<&'a HashMap<String, ConstExpression>>,
    /// Const fn evaluator for resolving function calls.
    pub evaluator: Option<&'a mut ConstEvaluator>,
}

/// Evaluates a const expression: substitutes bindings, folds arithmetic, resolves const fn calls.
pub fn eval(expression: &ConstExpression, context: &mut EvalContext) -> ConstExpression {
    use ConstExpression as C;
    match expression {
        C::Integer(_) => expression.clone(),

        C::Named(name) => {
            if let Some(bindings) = &context.bindings
                && let Some(value) = bindings.get(name)
            {
                return eval(value, context);
            }
            if let Some(evaluator) = &context.evaluator
                && let Some(value) = evaluator.resolve_named(name)
            {
                return C::Integer(value);
            }
            expression.clone()
        }

        C::Call(name, arguments) => {
            let resolved: Vec<_> = arguments.iter().map(|a| eval(a, context)).collect();
            if let Some(evaluator) = &mut context.evaluator
                && evaluator.has_function(name)
                && resolved.iter().all(|a| matches!(a, C::Integer(_)))
            {
                let values: Vec<_> = resolved
                    .iter()
                    .map(|a| match a {
                        C::Integer(v) => crate::const_eval::ConstValue::Integer(*v as i128),
                        _ => unreachable!(),
                    })
                    .collect();
                if let Ok(result) = evaluator.call(name, values)
                    && let Ok(const_expression) = result.to_const_expression()
                {
                    return const_expression;
                }
            }
            C::Call(name.clone(), resolved)
        }

        C::Field(inner, field) => {
            let resolved = eval(inner, context);
            if let C::Construction(fields) = &resolved
                && let Some((_, value)) = fields.iter().find(|(name, _)| name == field)
            {
                return value.clone();
            }
            C::Field(Box::new(resolved), field.clone())
        }

        C::Construction(fields) => C::Construction(
            fields
                .iter()
                .map(|(name, value)| (name.clone(), eval(value, context)))
                .collect(),
        ),

        binary => {
            let (left, right) = extract_binary_operands(binary);
            let left = eval(left, context);
            let right = eval(right, context);
            let folded = fold_binary(&left, &right, binary);
            folded.unwrap_or_else(|| reconstruct_binary(binary, left, right))
        }
    }
}

/// Evaluates a const expression and extracts the integer value.
pub fn eval_to_usize(expression: ConstExpression) -> Result<usize, String> {
    match expression {
        ConstExpression::Integer(value) => Ok(value as usize),
        other => Err(format!("expected integer constant, got {other:?}")),
    }
}

fn extract_binary_operands(expression: &ConstExpression) -> (&ConstExpression, &ConstExpression) {
    use ConstExpression as C;
    match expression {
        C::Add(l, r)
        | C::Subtract(l, r)
        | C::Multiply(l, r)
        | C::Divide(l, r)
        | C::Modulo(l, r)
        | C::BitwiseAnd(l, r)
        | C::BitwiseOr(l, r)
        | C::BitwiseXor(l, r)
        | C::ShiftLeft(l, r)
        | C::ShiftRight(l, r) => (l, r),
        _ => unreachable!(),
    }
}

fn fold_binary(
    left: &ConstExpression,
    right: &ConstExpression,
    original: &ConstExpression,
) -> Option<ConstExpression> {
    use ConstExpression as C;
    let (a, b) = match (left, right) {
        (C::Integer(a), C::Integer(b)) => (*a, *b),
        _ => return None,
    };
    let result = match original {
        C::Add(..) => a.wrapping_add(b),
        C::Subtract(..) => a.wrapping_sub(b),
        C::Multiply(..) => a.wrapping_mul(b),
        C::Divide(..) => {
            if b == 0 {
                return None;
            }
            a / b
        }
        C::Modulo(..) => {
            if b == 0 {
                return None;
            }
            a % b
        }
        C::BitwiseAnd(..) => a & b,
        C::BitwiseOr(..) => a | b,
        C::BitwiseXor(..) => a ^ b,
        C::ShiftLeft(..) => a.wrapping_shl(b as u32),
        C::ShiftRight(..) => a.wrapping_shr(b as u32),
        _ => unreachable!(),
    };
    Some(C::Integer(result))
}

fn reconstruct_binary(
    original: &ConstExpression,
    left: ConstExpression,
    right: ConstExpression,
) -> ConstExpression {
    use ConstExpression as C;
    let (left, right) = (Box::new(left), Box::new(right));
    match original {
        C::Add(..) => C::Add(left, right),
        C::Subtract(..) => C::Subtract(left, right),
        C::Multiply(..) => C::Multiply(left, right),
        C::Divide(..) => C::Divide(left, right),
        C::Modulo(..) => C::Modulo(left, right),
        C::BitwiseAnd(..) => C::BitwiseAnd(left, right),
        C::BitwiseOr(..) => C::BitwiseOr(left, right),
        C::BitwiseXor(..) => C::BitwiseXor(left, right),
        C::ShiftLeft(..) => C::ShiftLeft(left, right),
        C::ShiftRight(..) => C::ShiftRight(left, right),
        _ => unreachable!(),
    }
}
