//! Tree-walking interpreter for compile-time evaluation of `const fn`.

use std::collections::HashMap;

use t_ree::operator::{
    ArithmeticOperator, BinaryOperator, BitwiseOperator, ComparisonOperator, LogicalOperator,
    UnaryOperator,
};
use ty_ree::declaration as ty;
use ty_ree::expression::{self as ty_expression, ExpressionKind, Statement};
use ty_ree::types as ty_types;

const MAX_ITERATIONS: u64 = 100_000;

/// A compile-time constant value.
#[derive(Clone, Debug)]
pub enum ConstValue {
    /// Integer value.
    Integer(i128),
    /// Floating-point value.
    Float(f64),
    /// Boolean value.
    Bool(bool),
    /// Array of constant values.
    Array(Vec<Self>),
}

impl ConstValue {
    /// Extracts an integer value or returns an error.
    pub fn as_integer(&self) -> Result<i128, String> {
        match self {
            Self::Integer(value) => Ok(*value),
            other => Err(format!("expected integer, got {other:?}")),
        }
    }

    fn as_bool(&self) -> Result<bool, String> {
        match self {
            Self::Bool(value) => Ok(*value),
            other => Err(format!("expected bool, got {other:?}")),
        }
    }

    /// Converts to a const expression for use in type-level computations.
    pub fn to_const_expression(&self) -> Result<ty_types::ConstExpression, String> {
        match self {
            Self::Integer(value) => Ok(ty_types::ConstExpression::Integer(*value as u128)),
            other => Err(format!("cannot convert {other:?} to const expression")),
        }
    }

    /// Converts to a runtime AST expression.
    pub fn to_expression(&self) -> ty_expression::Expression {
        let (kind, resolved_type) = match self {
            Self::Integer(value) => (
                ExpressionKind::Literal(ty_expression::Literal::Integer(*value as u128)),
                ty_types::Type::Int(ty_types::IntWidth::W64, ty_types::Signedness::Signed),
            ),
            Self::Float(value) => (
                ExpressionKind::Literal(ty_expression::Literal::Float(*value)),
                ty_types::Type::Float(ty_types::FloatWidth::W64),
            ),
            Self::Bool(value) => (
                ExpressionKind::Literal(ty_expression::Literal::Bool(*value)),
                ty_types::Type::Bool,
            ),
            Self::Array(elements) => {
                let expressions: Vec<_> = elements.iter().map(Self::to_expression).collect();
                (
                    ExpressionKind::ArrayLiteral(expressions),
                    ty_types::Type::Inferred,
                )
            }
        };
        ty_expression::Expression::new(kind, resolved_type)
    }
}

/// A stored const function definition for compile-time evaluation.
#[derive(Clone, Debug)]
struct ConstFunction {
    /// Parameter names.
    parameter_names: Vec<String>,
    /// Function body.
    body: ty_expression::Block,
}

/// Tree-walking interpreter for const fn evaluation.
#[derive(Default)]
pub struct ConstEvaluator {
    const_functions: HashMap<String, ConstFunction>,
    constants: HashMap<String, u128>,
    scopes: Vec<HashMap<String, ConstValue>>,
    iteration_counter: u64,
}

impl ConstEvaluator {
    /// Creates a new empty evaluator.
    pub fn new() -> Self {
        Self::default()
    }

    /// Registers a const impl for compile-time evaluation, extracting
    /// parameter names from the type's inner structure.
    pub fn register_const_impl(
        &mut self,
        name: &str,
        newtypes: &HashMap<String, ty_types::Type>,
        impl_block: &ty::ImplBlock,
    ) {
        let parameter_names = if let Some(params) = &impl_block.callable.parameters {
            params.iter().map(|p| p.name.clone()).collect()
        } else if let Some(inner) = newtypes.get(name) {
            match inner {
                ty_types::Type::Named(field_name, _) => vec![field_name.clone()],
                ty_types::Type::Tuple(fields) => fields
                    .iter()
                    .map(|field| {
                        if let ty_types::Type::Named(field_name, _) = field {
                            field_name.clone()
                        } else {
                            "it".into()
                        }
                    })
                    .collect(),
                _ => vec!["it".into()],
            }
        } else {
            vec!["it".into()]
        };

        self.const_functions.insert(
            name.into(),
            ConstFunction {
                parameter_names,
                body: impl_block.callable.body.clone(),
            },
        );
    }

    /// Registers a named constant value.
    pub fn register_constant(&mut self, name: &str, value: u128) {
        self.constants.insert(name.into(), value);
    }

    /// Looks up a named constant by name.
    pub fn resolve_named(&self, name: &str) -> Option<u128> {
        self.constants.get(name).copied()
    }

    /// Returns true if any const functions or constants are registered.
    pub fn has_any(&self) -> bool {
        !self.const_functions.is_empty() || !self.constants.is_empty()
    }

    /// Returns true if a const fn with the given name is registered.
    pub fn has_function(&self, name: &str) -> bool {
        self.const_functions.contains_key(name)
    }

    /// Calls a registered const fn with the given arguments.
    pub fn call(&mut self, name: &str, arguments: Vec<ConstValue>) -> Result<ConstValue, String> {
        let function = self
            .const_functions
            .get(name)
            .ok_or_else(|| format!("undefined const function '{name}'"))?
            .clone();

        if arguments.len() != function.parameter_names.len() {
            return Err(format!(
                "const fn '{name}' expects {} arguments, got {}",
                function.parameter_names.len(),
                arguments.len()
            ));
        }

        let mut scope = HashMap::new();
        for (parameter_name, argument) in function.parameter_names.iter().zip(arguments) {
            scope.insert(parameter_name.clone(), argument);
        }
        self.scopes.push(scope);

        let result = self.evaluate_block(&function.body);

        self.scopes.pop();
        result
    }

    /// Evaluates a const expression to a concrete value.
    pub fn evaluate_const_expression(
        &mut self,
        expression: &ty_types::ConstExpression,
    ) -> Result<ConstValue, String> {
        use ty_types::ConstExpression as C;
        match expression {
            C::Integer(value) => Ok(ConstValue::Integer(*value as i128)),
            C::Named(name) => self.lookup(name),
            C::Call(name, arguments) => {
                let evaluated_arguments: Vec<_> = arguments
                    .iter()
                    .map(|argument| self.evaluate_const_expression(argument))
                    .collect::<Result<_, _>>()?;
                self.call(name, evaluated_arguments)
            }
            binary => {
                let (left, right, op) = match binary {
                    C::Add(l, r) => (l, r, "+"),
                    C::Subtract(l, r) => (l, r, "-"),
                    C::Multiply(l, r) => (l, r, "*"),
                    C::Divide(l, r) => (l, r, "/"),
                    C::Modulo(l, r) => (l, r, "%"),
                    C::BitwiseAnd(l, r) => (l, r, "&"),
                    C::BitwiseOr(l, r) => (l, r, "|"),
                    C::BitwiseXor(l, r) => (l, r, "^"),
                    C::ShiftLeft(l, r) => (l, r, "<<"),
                    C::ShiftRight(l, r) => (l, r, ">>"),
                    _ => unreachable!(),
                };
                let left_value = self.evaluate_const_expression(left)?.as_integer()?;
                let right_value = self.evaluate_const_expression(right)?.as_integer()?;
                let result = match op {
                    "+" => left_value.wrapping_add(right_value),
                    "-" => left_value.wrapping_sub(right_value),
                    "*" => left_value.wrapping_mul(right_value),
                    "/" => {
                        if right_value == 0 {
                            return Err("division by zero in const expression".into());
                        }
                        left_value / right_value
                    }
                    "%" => {
                        if right_value == 0 {
                            return Err("modulo by zero in const expression".into());
                        }
                        left_value % right_value
                    }
                    "&" => left_value & right_value,
                    "|" => left_value | right_value,
                    "^" => left_value ^ right_value,
                    "<<" => left_value.wrapping_shl(right_value as u32),
                    ">>" => left_value.wrapping_shr(right_value as u32),
                    _ => unreachable!(),
                };
                Ok(ConstValue::Integer(result))
            }
        }
    }

    fn evaluate_block(&mut self, block: &ty_expression::Block) -> Result<ConstValue, String> {
        match self.execute_block(block)? {
            Flow::Continue(value) => Ok(value),
            Flow::Return(value) => Ok(value),
        }
    }

    fn execute_block(&mut self, block: &ty_expression::Block) -> Result<Flow, String> {
        for statement in &block.statements {
            match self.execute_statement(statement)? {
                Flow::Continue(_) => {}
                flow @ Flow::Return(_) => return Ok(flow),
            }
        }
        let value = if let Some(result) = &block.result {
            self.evaluate_expression(result)?
        } else {
            ConstValue::Integer(0)
        };
        Ok(Flow::Continue(value))
    }

    fn execute_statement(&mut self, statement: &Statement) -> Result<Flow, String> {
        match statement {
            Statement::Expression(expression) => {
                self.evaluate_expression(expression)?;
                Ok(Flow::Continue(ConstValue::Integer(0)))
            }
            Statement::Let { name, value, .. } => {
                let value = self.evaluate_expression(value)?;
                self.define(name.clone(), value);
                Ok(Flow::Continue(ConstValue::Integer(0)))
            }
            Statement::Assign(target, value) => {
                let value = self.evaluate_expression(value)?;
                self.assign(target, value)?;
                Ok(Flow::Continue(ConstValue::Integer(0)))
            }
            Statement::Return(Some(expression)) => {
                let value = self.evaluate_expression(expression)?;
                Ok(Flow::Return(value))
            }
            Statement::Return(None) => Ok(Flow::Return(ConstValue::Integer(0))),
            Statement::While { condition, body } => {
                loop {
                    self.iteration_counter += 1;
                    if self.iteration_counter > MAX_ITERATIONS {
                        return Err(format!(
                            "const evaluation exceeded {MAX_ITERATIONS} iterations"
                        ));
                    }
                    let condition_value = self.evaluate_expression(condition)?.as_bool()?;
                    if !condition_value {
                        break;
                    }
                    match self.execute_block(body)? {
                        Flow::Continue(_) => {}
                        flow @ Flow::Return(_) => return Ok(flow),
                    }
                }
                Ok(Flow::Continue(ConstValue::Integer(0)))
            }
            other => Err(format!("unsupported statement in const context: {other:?}")),
        }
    }

    fn evaluate_expression(
        &mut self,
        expression: &ty_expression::Expression,
    ) -> Result<ConstValue, String> {
        match &expression.kind {
            ExpressionKind::Literal(literal) => Self::evaluate_literal(literal),
            ExpressionKind::Variable(name) => self.lookup(name),
            ExpressionKind::BinaryOperation(op, left, right) => {
                self.evaluate_binary(*op, left, right)
            }
            ExpressionKind::UnaryOperation(op, operand) => self.evaluate_unary(*op, operand),
            ExpressionKind::Call(callee, arguments) => {
                if let ExpressionKind::Variable(name) = &callee.kind {
                    let evaluated_arguments: Vec<_> = arguments
                        .iter()
                        .map(|argument| self.evaluate_expression(argument))
                        .collect::<Result<_, _>>()?;
                    self.call(name, evaluated_arguments)
                } else {
                    Err("only direct function calls supported in const context".into())
                }
            }
            ExpressionKind::Index(array, index) => {
                let array_value = self.evaluate_expression(array)?;
                let index_value = self.evaluate_expression(index)?.as_integer()?;
                match array_value {
                    ConstValue::Array(elements) => {
                        let index = index_value as usize;
                        elements
                            .get(index)
                            .cloned()
                            .ok_or_else(|| format!("index {index} out of bounds"))
                    }
                    _ => Err("indexing non-array in const context".into()),
                }
            }
            ExpressionKind::ArrayLiteral(elements) => {
                let values: Vec<_> = elements
                    .iter()
                    .map(|element| self.evaluate_expression(element))
                    .collect::<Result<_, _>>()?;
                Ok(ConstValue::Array(values))
            }
            ExpressionKind::Block(block) => self.evaluate_block(block),
            ExpressionKind::If {
                condition,
                then_branch,
                else_branch,
            } => {
                let condition_value = self.evaluate_expression(condition)?.as_bool()?;
                if condition_value {
                    self.evaluate_block(then_branch)
                } else if let Some(else_branch) = else_branch {
                    self.evaluate_block(else_branch)
                } else {
                    Ok(ConstValue::Integer(0))
                }
            }
            ExpressionKind::Convert(inner, target_type) => {
                let value = self.evaluate_expression(inner)?;
                Self::evaluate_cast(value, target_type)
            }
            ExpressionKind::Field(inner, _) => self.evaluate_expression(inner),
            ExpressionKind::Replace(target, value) => {
                let new_value = self.evaluate_expression(value)?;
                let old_value = self.evaluate_place(target)?;
                self.assign(target, new_value)?;
                Ok(old_value)
            }
            ExpressionKind::Dereference(inner) => self.evaluate_expression(inner),
            other => Err(format!(
                "unsupported expression in const context: {other:?}"
            )),
        }
    }

    fn evaluate_literal(literal: &ty_expression::Literal) -> Result<ConstValue, String> {
        match literal {
            ty_expression::Literal::Integer(value) => Ok(ConstValue::Integer(*value as i128)),
            ty_expression::Literal::Float(value) => Ok(ConstValue::Float(*value)),
            ty_expression::Literal::Bool(value) => Ok(ConstValue::Bool(*value)),
            ty_expression::Literal::String(_) => {
                Err("string literals not supported in const context".into())
            }
            ty_expression::Literal::Null => Err("null not supported in const context".into()),
        }
    }

    fn evaluate_binary(
        &mut self,
        op: BinaryOperator,
        left: &ty_expression::Expression,
        right: &ty_expression::Expression,
    ) -> Result<ConstValue, String> {
        let left_value = self.evaluate_expression(left)?;
        let right_value = self.evaluate_expression(right)?;

        match op {
            BinaryOperator::Arithmetic(arithmetic) => {
                let (a, b) = (left_value.as_integer()?, right_value.as_integer()?);
                let result = match arithmetic {
                    ArithmeticOperator::Add => a.wrapping_add(b),
                    ArithmeticOperator::Subtract => a.wrapping_sub(b),
                    ArithmeticOperator::Multiply => a.wrapping_mul(b),
                    ArithmeticOperator::Divide => {
                        if b == 0 {
                            return Err("division by zero".into());
                        }
                        a / b
                    }
                    ArithmeticOperator::Remainder => {
                        if b == 0 {
                            return Err("remainder by zero".into());
                        }
                        a % b
                    }
                };
                Ok(ConstValue::Integer(result))
            }
            BinaryOperator::Comparison(comparison) => {
                let (a, b) = (left_value.as_integer()?, right_value.as_integer()?);
                let result = match comparison {
                    ComparisonOperator::Equal => a == b,
                    ComparisonOperator::NotEqual => a != b,
                    ComparisonOperator::Less => a < b,
                    ComparisonOperator::LessEqual => a <= b,
                    ComparisonOperator::Greater => a > b,
                    ComparisonOperator::GreaterEqual => a >= b,
                };
                Ok(ConstValue::Bool(result))
            }
            BinaryOperator::Logical(logical) => {
                let (a, b) = (left_value.as_bool()?, right_value.as_bool()?);
                let result = match logical {
                    LogicalOperator::And => a && b,
                    LogicalOperator::Or => a || b,
                };
                Ok(ConstValue::Bool(result))
            }
            BinaryOperator::Bitwise(bitwise) => {
                let (a, b) = (left_value.as_integer()?, right_value.as_integer()?);
                let result = match bitwise {
                    BitwiseOperator::And => a & b,
                    BitwiseOperator::Or => a | b,
                    BitwiseOperator::Xor => a ^ b,
                    BitwiseOperator::ShiftLeft => a.wrapping_shl(b as u32),
                    BitwiseOperator::ShiftRight => a.wrapping_shr(b as u32),
                };
                Ok(ConstValue::Integer(result))
            }
        }
    }

    fn evaluate_unary(
        &mut self,
        op: UnaryOperator,
        operand: &ty_expression::Expression,
    ) -> Result<ConstValue, String> {
        let value = self.evaluate_expression(operand)?;
        match op {
            UnaryOperator::Negate => Ok(ConstValue::Integer(-value.as_integer()?)),
            UnaryOperator::LogicalNot => Ok(ConstValue::Bool(!value.as_bool()?)),
            UnaryOperator::BitwiseNot => Ok(ConstValue::Integer(!value.as_integer()?)),
        }
    }

    fn evaluate_cast(
        value: ConstValue,
        target_type: &ty_types::Type,
    ) -> Result<ConstValue, String> {
        match (value, target_type) {
            (ConstValue::Integer(value), ty_types::Type::Float(..)) => {
                Ok(ConstValue::Float(value as f64))
            }
            (ConstValue::Float(value), ty_types::Type::Int(..)) => {
                Ok(ConstValue::Integer(value as i128))
            }
            (value @ ConstValue::Integer(_), ty_types::Type::Int(..)) => Ok(value),
            (value @ ConstValue::Float(_), ty_types::Type::Float(..)) => Ok(value),
            (value, ty_types::Type::Bool) => match value {
                ConstValue::Integer(value) => Ok(ConstValue::Bool(value != 0)),
                ConstValue::Bool(value) => Ok(ConstValue::Bool(value)),
                other => Err(format!("cannot cast {other:?} to bool")),
            },
            (value, _) => Err(format!("unsupported cast in const context: {value:?}")),
        }
    }

    fn lookup(&self, name: &str) -> Result<ConstValue, String> {
        for scope in self.scopes.iter().rev() {
            if let Some(value) = scope.get(name) {
                return Ok(value.clone());
            }
        }
        Err(format!("undefined variable '{name}' in const context"))
    }

    fn define(&mut self, name: String, value: ConstValue) {
        if let Some(scope) = self.scopes.last_mut() {
            scope.insert(name, value);
        }
    }

    fn assign(
        &mut self,
        target: &ty_expression::Expression,
        value: ConstValue,
    ) -> Result<(), String> {
        match &target.kind {
            ExpressionKind::Variable(name) => {
                for scope in self.scopes.iter_mut().rev() {
                    if scope.contains_key(name) {
                        scope.insert(name.clone(), value);
                        return Ok(());
                    }
                }
                Err(format!("undefined variable '{name}' in const assignment"))
            }
            ExpressionKind::Dereference(inner) => self.assign(inner, value),
            ExpressionKind::Index(array, index) => {
                let index_value = self.evaluate_expression(index)?.as_integer()? as usize;
                let array = match &array.kind {
                    ExpressionKind::Variable(name) => name,
                    ExpressionKind::Dereference(inner) => {
                        if let ExpressionKind::Variable(name) = &inner.kind {
                            name
                        } else {
                            return Err(
                                "only simple array index assignment supported in const context"
                                    .into(),
                            );
                        }
                    }
                    _ => {
                        return Err(
                            "only simple array index assignment supported in const context".into(),
                        );
                    }
                };
                for scope in self.scopes.iter_mut().rev() {
                    if let Some(ConstValue::Array(elements)) = scope.get_mut(array) {
                        if index_value < elements.len() {
                            elements[index_value] = value;
                            return Ok(());
                        }
                        return Err(format!("index {index_value} out of bounds"));
                    }
                }
                Err(format!("undefined array '{array}' in const assignment"))
            }
            other => Err(format!(
                "unsupported assignment target in const context: {other:?}"
            )),
        }
    }

    fn evaluate_place(
        &mut self,
        expression: &ty_expression::Expression,
    ) -> Result<ConstValue, String> {
        self.evaluate_expression(expression)
    }
}

enum Flow {
    Continue(ConstValue),
    Return(ConstValue),
}
