rlox/src/backend/interpreter.rs

use crate::{
    backend::environment::EnvironmentStack,
    frontend::{
        ast::{AstNode, AstNodeKind, Expr, Stmt},
        source_registry::SourceSlice,
        tokens::{LiteralValue, NativeFunction, TokenType},
    },
    logging::display_ast::PrettyPrintExt,
    result::{runtime_error, LoxError, LoxResult},
};
use std::{
    fmt::{Debug, Display},
    ops::{Add, Div, Mul, Neg, Not, Rem, Sub},
};

fn error(message: String) -> LoxError {
    LoxError::RuntimeError {
        source_slice: SourceSlice::default(), // todo change this with the actual source slice
        message: message,
    }
}

fn error_at(source_slice: SourceSlice, message: String) -> LoxError {
    LoxError::RuntimeError {
        source_slice,
        message,
    }
}

impl Not for LiteralValue {
    type Output = LiteralValue;

    fn not(self) -> Self::Output {
        match self {
            LiteralValue::Boolean(b) => LiteralValue::Boolean(!b),
            LiteralValue::Number(n) => LiteralValue::Boolean(n == 0.0),
            _ => LiteralValue::Boolean(false),
        }
    }
}

pub trait Truthy {
    fn is_truthy(&self) -> bool;
}

impl Truthy for LiteralValue {
    fn is_truthy(&self) -> bool {
        match self {
            LiteralValue::Boolean(b) => *b,
            LiteralValue::Number(n) => *n != 0.0,
            LiteralValue::String(s) => !s.is_empty(),
            _ => false,
        }
    }
}

impl Neg for LiteralValue {
    type Output = LoxResult<LiteralValue>;

    fn neg(self) -> Self::Output {
        Ok(LiteralValue::Boolean(!self.is_truthy()))
    }
}

impl Add for LiteralValue {
    type Output = LoxResult<LiteralValue>;

    fn add(self, other: LiteralValue) -> Self::Output {
        match (self, other) {
            (LiteralValue::Number(a), LiteralValue::Number(b)) => Ok(LiteralValue::Number(a + b)),
            (LiteralValue::String(a), LiteralValue::String(b)) => {
                Ok(LiteralValue::String(format!("{}{}", a, b)))
            }
            _ => Err(error("Cannot add non-numeric values".to_string())),
        }
    }
}

impl LiteralValue {
    pub fn add_with_source(
        self,
        other: LiteralValue,
        source_slice: SourceSlice,
    ) -> LoxResult<LiteralValue> {
        match (self, other) {
            (LiteralValue::Number(a), LiteralValue::Number(b)) => Ok(LiteralValue::Number(a + b)),
            (LiteralValue::String(a), LiteralValue::String(b)) => {
                Ok(LiteralValue::String(format!("{}{}", a, b)))
            }
            _ => Err(error_at(
                source_slice,
                "Cannot add non-numeric values".to_string(),
            )),
        }
    }
}

impl Sub for LiteralValue {
    type Output = LoxResult<LiteralValue>;

    fn sub(self, other: LiteralValue) -> Self::Output {
        match (self, other) {
            (LiteralValue::Number(a), LiteralValue::Number(b)) => Ok(LiteralValue::Number(a - b)),
            _ => Err(error("Cannot subtract non-numeric values".to_string())),
        }
    }
}

impl Div for LiteralValue {
    type Output = LoxResult<LiteralValue>;

    fn div(self, other: LiteralValue) -> Self::Output {
        match (self, other) {
            (LiteralValue::Number(a), LiteralValue::Number(b)) => {
                if b == 0.0 {
                    Err(error("Division by zero".to_string()))
                } else {
                    Ok(LiteralValue::Number(a / b))
                }
            }
            _ => Err(error("Cannot divide non-numeric values".to_string())),
        }
    }
}

impl Mul for LiteralValue {
    type Output = LoxResult<LiteralValue>;

    fn mul(self, other: LiteralValue) -> Self::Output {
        match (self, other) {
            (LiteralValue::Number(a), LiteralValue::Number(b)) => Ok(LiteralValue::Number(a * b)),
            _ => Err(error("Cannot multiply non-numeric values".to_string())),
        }
    }
}

impl Rem for LiteralValue {
    type Output = LoxResult<LiteralValue>;

    fn rem(self, other: LiteralValue) -> Self::Output {
        match (self, other) {
            (LiteralValue::Number(a), LiteralValue::Number(b)) => {
                if b == 0.0 {
                    Err(error("Division by zero".to_string()))
                } else {
                    Ok(LiteralValue::Number(a % b))
                }
            }
            _ => Err(error("Cannot divide non-numeric values".to_string())),
        }
    }
}

impl PartialOrd for LiteralValue {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        match (self, other) {
            (LiteralValue::Number(a), LiteralValue::Number(b)) => a.partial_cmp(b),
            _ => None,
        }
    }
}

impl From<LiteralValue> for bool {
    fn from(value: LiteralValue) -> Self {
        match value {
            LiteralValue::Boolean(b) => b,
            LiteralValue::Number(n) => n != 0.0,
            _ => false,
        }
    }
}

pub struct Interpreter {
    enviorment: EnvironmentStack,
}

pub trait EvaluateInterpreter<T> {
    fn evaluate(&mut self, stmt: T) -> LoxResult<LiteralValue>;
}

impl<'a, R: AstNodeKind + Clone + Debug + Display> EvaluateInterpreter<AstNode<R>> for Interpreter
where
    Interpreter: EvaluateInterpreter<R>,
{
    fn evaluate(&mut self, stmt: AstNode<R>) -> LoxResult<LiteralValue> {
        match self.evaluate(stmt.node.clone()) {
            Ok(value) => Ok(value),
            Err(err) => Err(LoxError::RuntimeError {
                source_slice: stmt.source_slice,
                message: err.get_message(),
            }),
        }
    }
}

// Direct Expr evaluation to avoid infinite recursion
impl EvaluateInterpreter<Expr> for Interpreter {
    fn evaluate(&mut self, expr: Expr) -> LoxResult<LiteralValue> {
        match expr {
            Expr::Literal { value } => Ok(value),
            Expr::Identifier { name } => self.enviorment.get(&name),
            Expr::Binary {
                left,
                operator,
                right,
            } => {
                let left_val = self.evaluate(*left)?;
                let right_val = self.evaluate(*right)?;
                self.evaluate_binary(left_val, operator, right_val)
            }
            Expr::Unary { operator, operand } => {
                let operand_val = self.evaluate(*operand)?;
                self.evaluate_unary(operator, operand_val)
            }
            Expr::Grouping { expression } => self.evaluate(*expression),
            Expr::Call { callee, arguments } => self.evaluate_call(callee, arguments),
        }
    }
}

impl EvaluateInterpreter<AstNode<Stmt>> for Interpreter {
    fn evaluate(&mut self, node: AstNode<Stmt>) -> LoxResult<LiteralValue> {
        let stmt = node.node;
        let result = self.interpret_stmt_inner(stmt);
        match result {
            Ok(value) => Ok(value),
            Err(LoxError::RuntimeError {
                message,
                source_slice,
            }) if source_slice == SourceSlice::default() => {
                runtime_error(node.source_slice.clone(), message)
            }
            Err(err) => Err(err),
        }
    }
}

impl Interpreter {
    pub fn new() -> Self {
        let mut env = EnvironmentStack::new();
        let _ = env.declare(
            "clock".to_string(),
            LiteralValue::NativeFunction(NativeFunction::new(0, |_args| {
                LiteralValue::Number(
                    std::time::SystemTime::now()
                        .duration_since(std::time::UNIX_EPOCH)
                        .unwrap()
                        .as_millis() as f64,
                )
            })),
        );
        Self { enviorment: env }
    }
    fn evaluate_call(
        &mut self,
        callee: Box<AstNode<Expr>>,
        arguments: Vec<AstNode<Expr>>,
    ) -> LoxResult<LiteralValue> {
        let source_slice = callee.source_slice.clone();

        // Estrai il nome della variabile se il callee è un identificatore
        let function_name = match &callee.node {
            Expr::Identifier { name } => Some(name.clone()),
            _ => None,
        };

        let function = if let Some(name) = function_name {
            // Se abbiamo un nome di variabile, ottieni la funzione dall'ambiente
            self.enviorment.get(&name)?
        } else {
            // Altrimenti valuta l'espressione callee (per casi più complessi)
            self.evaluate(*callee)?
        };

        if !function.is_callable() {
            return runtime_error(source_slice, "Can only call functions");
        }

        let evaluated_arguments = arguments
            .iter()
            .map(|arg| self.evaluate(arg.clone()))
            .collect::<Result<Vec<LiteralValue>, LoxError>>()?;

        match function {
            LiteralValue::NativeFunction(func) => Ok((func.function)(&evaluated_arguments)),
            LiteralValue::Function(func) => {
                func.parameters
                    .iter()
                    .enumerate()
                    .map(|(index, par)| {
                        let value = evaluated_arguments.get(index).unwrap();
                        self.enviorment.declare(par.0.clone(), value.clone())
                    })
                    .collect::<LoxResult<Vec<LiteralValue>>>()?;
                self.evaluate(func.body)
            }
            _ => runtime_error(
                source_slice.clone(),
                "This is callable, but apparently is not implemented",
            ),
        }
    }

    fn evaluate_binary(
        &mut self,
        left: LiteralValue,
        operator: TokenType,
        right: LiteralValue,
    ) -> LoxResult<LiteralValue> {
        match operator {
            TokenType::Plus => left + right,
            TokenType::Minus => left - right,
            TokenType::Star => left * right,
            TokenType::Slash => left / right,
            TokenType::Greater => Ok(LiteralValue::Boolean(left > right)),
            TokenType::GreaterEqual => Ok(LiteralValue::Boolean(left >= right)),
            TokenType::Less => Ok(LiteralValue::Boolean(left < right)),
            TokenType::LessEqual => Ok(LiteralValue::Boolean(left <= right)),
            TokenType::EqualEqual => Ok(LiteralValue::Boolean(left == right)),
            TokenType::BangEqual => Ok(LiteralValue::Boolean(left != right)),
            TokenType::And => Ok(LiteralValue::Boolean(left.is_truthy() && right.is_truthy())),
            TokenType::Or => Ok(LiteralValue::Boolean(left.is_truthy() || right.is_truthy())),
            _ => Err(error(format!(
                "Unsupported binary operator: {:?}",
                operator
            ))),
        }
    }

    fn evaluate_unary(
        &mut self,
        operator: TokenType,
        operand: LiteralValue,
    ) -> LoxResult<LiteralValue> {
        match operator {
            TokenType::Minus => -operand,
            TokenType::Bang => Ok(!operand),
            _ => Err(error(format!("Unsupported unary operator: {:?}", operator))),
        }
    }

    fn interpret_stmt_inner(&mut self, stmt: Stmt) -> LoxResult<LiteralValue> {
        match stmt {
            Stmt::Expression { expression } => self.evaluate(*expression),
            Stmt::Print { expression } => {
                let value = self.evaluate(*expression)?;
                println!("print interpreter: \t{}", value);
                Ok(LiteralValue::Nil)
            }
            Stmt::Block { statements } => self.evaluate_block(*statements),
            Stmt::Stmt { expression } => self.evaluate(*expression.clone()),
            Stmt::Return { expression } => self.evaluate(*expression),
            Stmt::VarDeclaration { name, initializer } => {
                let value = match initializer {
                    Some(expr_node) => match &expr_node.node {
                        Expr::Literal { value } => {
                            if value.is_callable() {
                                value.clone()
                            } else {
                                value.clone()
                            }
                        }
                        _ => self.evaluate(*expr_node)?,
                    },
                    None => LiteralValue::Nil,
                };
                self.enviorment.declare(name.clone(), value)
            }

            Stmt::VarAssigment { name, value } => {
                let result = self.evaluate(*value)?;
                self.enviorment.set(name.clone(), result)
            }
            Stmt::If {
                condition,
                then_branch,
                elif_branch,
                else_branch,
            } => self.evaluate_if(condition, then_branch, elif_branch, else_branch),
            Stmt::While { condition, body } => {
                while self.evaluate(*condition.clone())?.is_truthy() {
                    self.evaluate(*body.clone())?;
                }
                Ok(LiteralValue::Nil)
            }
            Stmt::For {
                variable,
                condition,
                increment,
                body,
            } => {
                self.enviorment.push_new_scope();
                let source_slice = variable.source_slice.clone();
                let val = self.evaluate(*variable)?;
                if !matches!(val, LiteralValue::Number(..)) {
                    return runtime_error(source_slice, "Expected number literal");
                }
                let mut ret = LiteralValue::Nil;
                while self.evaluate(*condition.clone())?.is_truthy() {
                    ret = self.evaluate(*body.clone())?;
                    self.evaluate(*increment.clone())?;
                }

                Ok(ret)
            }
        }
    }
    fn evaluate_if(
        &mut self,
        condition: Box<AstNode<Expr>>,
        then_branch: Box<AstNode<Stmt>>,
        elif_branch: Vec<(Box<AstNode<Expr>>, Box<AstNode<Stmt>>)>,
        else_branch: Option<Box<AstNode<Stmt>>>,
    ) -> LoxResult<LiteralValue> {
        let condition = self.evaluate(*condition)?;
        match condition {
            LiteralValue::Boolean(true) => self.evaluate(*then_branch),
            LiteralValue::Boolean(false) => {
                for (elif_condition, elif_then_branch) in elif_branch {
                    let condition = self.evaluate(*elif_condition)?;
                    match condition {
                        LiteralValue::Boolean(true) => {
                            return self.evaluate(*elif_then_branch);
                        }
                        LiteralValue::Boolean(false) => continue,
                        _ => {
                            return Err(LoxError::TypeMismatch {
                                source_slice: SourceSlice::default(), // todo change this to the actual source slice
                                expected: "boolean".to_string(),
                                found: condition.to_string(),
                            });
                        }
                    };
                }
                if let Some(else_block) = else_branch {
                    self.evaluate(*else_block)
                } else {
                    Ok(LiteralValue::Nil)
                }
            }
            _ => Err(LoxError::TypeMismatch {
                source_slice: SourceSlice::default(), // todo change this to the actual source slice
                expected: "boolean".to_string(),
                found: condition.to_string(),
            }),
        }
    }

    fn evaluate_block(&mut self, statements: Vec<AstNode<Stmt>>) -> LoxResult<LiteralValue> {
        self.enviorment.push_new_scope();
        let (elements, final_expr) = match statements.split_last() {
            Some((stmt, body)) => match &stmt.node {
                Stmt::Expression { expression } => (body, Some(expression)),
                Stmt::Return { expression } => (body, Some(expression)),
                _ => (statements.as_slice(), None),
            },

            None => {
                (&[][..], None) // Blocco vuoto
            }
        };

        // Ora elements è sempre disponibile
        for statement in elements.iter() {
            self.evaluate((*statement).clone())?;
        }

        // Gestisci l'espressione finale se presente
        match final_expr {
            Some(expr) => {
                let res = self.evaluate(*expr.clone());
                self.enviorment.pop_scope();
                res
            }
            None => {
                self.enviorment.pop_scope();
                Ok(LiteralValue::Nil)
            }
        }
    }
}