rlox/src/common/base_value.rs

use core::fmt;
use std::collections::HashMap;
use std::fmt::Display;
use std::ops::{Add, Div, Mul, Not, Rem, Sub};
use std::ptr::fn_addr_eq;

use crate::common::lox_result::{LoxResult, runtime_error};
use crate::common::types::Type;
use crate::{
    backend::environment::EnvironmentStack, common::ast::AstNode,
    frontend::source_registry::SourceSlice,
};
#[derive(Debug, Clone, PartialEq)]
pub enum Number {
    I16(i16),
    I32(i32),
    I64(i64),
    F32(f32),
    F64(f64),
    U8(u8),
    U16(u16),
    U32(u32),
    U64(u64),
    U128(u128),
}

impl Display for Number {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Number::I16(n) => write!(f, "{}", n),
            Number::I32(n) => write!(f, "{}", n),
            Number::I64(n) => write!(f, "{}", n),
            Number::F32(n) => write!(f, "{}", n),
            Number::F64(n) => write!(f, "{}", n),
            Number::U8(n) => write!(f, "{}", n),
            Number::U16(n) => write!(f, "{}", n),
            Number::U32(n) => write!(f, "{}", n),
            Number::U64(n) => write!(f, "{}", n),
            Number::U128(n) => write!(f, "{}", n),
        }
    }
}

impl Number {
    /// Promotes two numbers to their common type following the hierarchy:
    /// Unsigned -> Integer -> Float
    fn promote_to_common_type(a: Number, b: Number) -> (Number, Number) {
        use Number::*;
        match (&a, &b) {
            // If either is a float, both become F64
            (F32(_), _) | (F64(_), _) => {
                let a_as_f64 = match a {
                    I16(n) => n as f64,
                    I32(n) => n as f64,
                    I64(n) => n as f64,
                    F32(n) => n as f64,
                    F64(n) => n,
                    U8(n) => n as f64,
                    U16(n) => n as f64,
                    U32(n) => n as f64,
                    U64(n) => n as f64,
                    U128(n) => n as f64,
                };
                let b_as_f64 = match b {
                    I16(n) => n as f64,
                    I32(n) => n as f64,
                    I64(n) => n as f64,
                    F32(n) => n as f64,
                    F64(n) => n,
                    U8(n) => n as f64,
                    U16(n) => n as f64,
                    U32(n) => n as f64,
                    U64(n) => n as f64,
                    U128(n) => n as f64,
                };
                (F64(a_as_f64), F64(b_as_f64))
            }
            (_, F32(_)) | (_, F64(_)) => {
                // Swap and recurse to handle the float case
                let (b_promoted, a_promoted) = Self::promote_to_common_type(b.clone(), a.clone());
                (a_promoted, b_promoted)
            }
            // If either is signed integer, both become I64
            (I16(_), _) | (I32(_), _) | (I64(_), _) => {
                let a_as_i64 = match a {
                    I16(n) => n as i64,
                    I32(n) => n as i64,
                    I64(n) => n,
                    U8(n) => n as i64,
                    U16(n) => n as i64,
                    U32(n) => n as i64,
                    U64(n) => n as i64,
                    U128(n) => n as i64, // May overflow but keeping it simple
                    _ => unreachable!(),
                };
                let b_as_i64 = match b {
                    I16(n) => n as i64,
                    I32(n) => n as i64,
                    I64(n) => n,
                    U8(n) => n as i64,
                    U16(n) => n as i64,
                    U32(n) => n as i64,
                    U64(n) => n as i64,
                    U128(n) => n as i64, // May overflow but keeping it simple
                    _ => unreachable!(),
                };
                (I64(a_as_i64), I64(b_as_i64))
            }
            (_, I16(_)) | (_, I32(_)) | (_, I64(_)) => {
                // Swap and recurse to handle the signed case
                let (b_promoted, a_promoted) = Self::promote_to_common_type(b.clone(), a.clone());
                (a_promoted, b_promoted)
            }
            // Both are unsigned, promote to U128
            _ => {
                let a_as_u128 = match a {
                    U8(n) => n as u128,
                    U16(n) => n as u128,
                    U32(n) => n as u128,
                    U64(n) => n as u128,
                    U128(n) => n,
                    _ => unreachable!(),
                };
                let b_as_u128 = match b {
                    U8(n) => n as u128,
                    U16(n) => n as u128,
                    U32(n) => n as u128,
                    U64(n) => n as u128,
                    U128(n) => n,
                    _ => unreachable!(),
                };
                (U128(a_as_u128), U128(b_as_u128))
            }
        }
    }

    pub fn add(self, other: Number) -> Number {
        let (a, b) = Self::promote_to_common_type(self, other);
        match (a, b) {
            (Number::F64(x), Number::F64(y)) => Number::F64(x + y),
            (Number::I64(x), Number::I64(y)) => Number::I64(x + y),
            (Number::U128(x), Number::U128(y)) => Number::U128(x + y),
            _ => unreachable!("promote_to_common_type should handle all cases"),
        }
    }

    pub fn sub(self, other: Number) -> Number {
        let (a, b) = Self::promote_to_common_type(self, other);
        match (a, b) {
            (Number::F64(x), Number::F64(y)) => Number::F64(x - y),
            (Number::I64(x), Number::I64(y)) => Number::I64(x - y),
            (Number::U128(x), Number::U128(y)) => Number::U128(x - y),
            _ => unreachable!("promote_to_common_type should handle all cases"),
        }
    }

    pub fn mul(self, other: Number) -> Number {
        let (a, b) = Self::promote_to_common_type(self, other);
        match (a, b) {
            (Number::F64(x), Number::F64(y)) => Number::F64(x * y),
            (Number::I64(x), Number::I64(y)) => Number::I64(x * y),
            (Number::U128(x), Number::U128(y)) => Number::U128(x * y),
            _ => unreachable!("promote_to_common_type should handle all cases"),
        }
    }

    pub fn div(self, other: Number) -> Option<Number> {
        let (a, b) = Self::promote_to_common_type(self, other);
        match (a, b) {
            (Number::F64(x), Number::F64(y)) => {
                if y == 0.0 {
                    None
                } else {
                    Some(Number::F64(x / y))
                }
            }
            (Number::I64(x), Number::I64(y)) => {
                if y == 0 {
                    None
                } else {
                    Some(Number::I64(x / y))
                }
            }
            (Number::U128(x), Number::U128(y)) => {
                if y == 0 {
                    None
                } else {
                    Some(Number::U128(x / y))
                }
            }
            _ => unreachable!("promote_to_common_type should handle all cases"),
        }
    }

    pub fn rem(self, other: Number) -> Option<Number> {
        let (a, b) = Self::promote_to_common_type(self, other);
        match (a, b) {
            (Number::F64(x), Number::F64(y)) => {
                if y == 0.0 {
                    None
                } else {
                    Some(Number::F64(x % y))
                }
            }
            (Number::I64(x), Number::I64(y)) => {
                if y == 0 {
                    None
                } else {
                    Some(Number::I64(x % y))
                }
            }
            (Number::U128(x), Number::U128(y)) => {
                if y == 0 {
                    None
                } else {
                    Some(Number::U128(x % y))
                }
            }
            _ => unreachable!("promote_to_common_type should handle all cases"),
        }
    }

    pub fn neg(self) -> Number {
        match self {
            Number::I16(n) => Number::I16(-n),
            Number::I32(n) => Number::I32(-n),
            Number::I64(n) => Number::I64(-n),
            Number::F32(n) => Number::F32(-n),
            Number::F64(n) => Number::F64(-n),
            // Unsigned numbers become signed when negated
            Number::U8(n) => Number::I32(-(n as i32)),
            Number::U16(n) => Number::I32(-(n as i32)),
            Number::U32(n) => Number::I64(-(n as i64)),
            Number::U64(n) => Number::I64(-(n as i64)),
            Number::U128(n) => Number::I64(-(n as i64)), // May overflow
        }
    }

    pub fn is_zero(&self) -> bool {
        match self {
            Number::I16(n) => *n == 0,
            Number::I32(n) => *n == 0,
            Number::I64(n) => *n == 0,
            Number::F32(n) => *n == 0.0,
            Number::F64(n) => *n == 0.0,
            Number::U8(n) => *n == 0,
            Number::U16(n) => *n == 0,
            Number::U32(n) => *n == 0,
            Number::U64(n) => *n == 0,
            Number::U128(n) => *n == 0,
        }
    }
}

impl PartialOrd for Number {
    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
        let (a, b) = Self::promote_to_common_type(self.clone(), other.clone());
        match (a, b) {
            (Number::F64(x), Number::F64(y)) => x.partial_cmp(&y),
            (Number::I64(x), Number::I64(y)) => x.partial_cmp(&y),
            (Number::U128(x), Number::U128(y)) => x.partial_cmp(&y),
            _ => unreachable!(),
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub enum BaseValue {
    Identifier(String),
    String(String),
    Number(Number),
    Boolean(bool),
    Nil,
    Function(Box<LoxFunction>),
    NativeFunction(Box<NativeFunction>),
    Struct(Box<Struct>),
}

struct Dict {
    map: HashMap<String, BaseValue>,
}

struct ReturnValue {
    label: Vec<String>,
    value: Type,
}

impl BaseValue {
    pub fn is_callable(&self) -> bool {
        matches!(
            self,
            BaseValue::Function(..) | BaseValue::NativeFunction(..)
        )
    }
}

impl Display for BaseValue {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            BaseValue::Identifier(id) => write!(f, "{}", id),
            BaseValue::String(s) => write!(f, "{}", s),
            BaseValue::Number(n) => write!(f, "{}", n),
            BaseValue::Boolean(b) => write!(f, "{}", b),
            BaseValue::Nil => write!(f, "nil"),
            BaseValue::Function(..) => write!(f, "<fn lox function>"),
            BaseValue::NativeFunction(..) => write!(f, "<native native function>"),
            BaseValue::Struct(_) => write!(f, "<struct>"),
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct LoxFunction {
    pub parameters: Vec<(String, Type)>,
    pub return_type: Type,
    pub body: AstNode,
    pub closure: Option<EnvironmentStack<BaseValue>>,
    pub guard: Option<Box<AstNode>>,
}

impl LoxFunction {
    pub fn new(
        parameters: Vec<(String, Type)>,
        body: AstNode,
        closure: Option<EnvironmentStack<BaseValue>>,
        guard: Option<Box<AstNode>>,
    ) -> Self {
        LoxFunction {
            parameters,
            return_type: Type::Any,
            body,
            closure,
            guard,
        }
    }

    pub fn anonymous_function(parameters: Vec<(String, Type)>, body: AstNode) -> Self {
        LoxFunction {
            parameters,
            return_type: Type::Any,
            body,
            closure: None,
            guard: None,
        }
    }
}

#[derive(Debug, Clone)]
pub struct NativeFunction {
    pub parameters: Vec<(String, Type)>,
    pub function: fn(&[BaseValue]) -> BaseValue,
}
impl PartialEq for NativeFunction {
    fn eq(&self, other: &Self) -> bool {
        self.parameters == other.parameters && fn_addr_eq(self.function, other.function)
    }
}

impl NativeFunction {
    pub fn new(parameters: Vec<(String, Type)>, function: fn(&[BaseValue]) -> BaseValue) -> Self {
        NativeFunction {
            parameters,
            function,
        }
    }
}

#[derive(Debug, Clone, PartialEq)]
pub struct Struct {
    pub fields: Vec<(String, Type)>,
}

// Trait implementations for BaseValue

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

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

impl Not for BaseValue {
    type Output = BaseValue;

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

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

    fn add(self, other: BaseValue) -> Self::Output {
        match (self, other) {
            (BaseValue::Number(a), BaseValue::Number(b)) => Ok(BaseValue::Number(a.add(b))),
            (BaseValue::String(a), BaseValue::String(b)) => {
                Ok(BaseValue::String(format!("{}{}", a, b)))
            }
            _ => runtime_error(
                SourceSlice::synthetic(),
                "Cannot add non-numeric values".to_string(),
            ),
        }
    }
}

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

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

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

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

    fn div(self, other: BaseValue) -> Self::Output {
        match (self, other) {
            (BaseValue::Number(a), BaseValue::Number(b)) => match a.div(b) {
                Some(result) => Ok(BaseValue::Number(result)),
                None => runtime_error(SourceSlice::synthetic(), "Division by zero".to_string()),
            },
            _ => runtime_error(
                SourceSlice::synthetic(),
                "Cannot divide non-numeric values".to_string(),
            ),
        }
    }
}

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

    fn mul(self, other: BaseValue) -> Self::Output {
        match (self, other) {
            (BaseValue::Number(a), BaseValue::Number(b)) => Ok(BaseValue::Number(a.mul(b))),
            _ => runtime_error(
                SourceSlice::synthetic(),
                "Cannot multiply non-numeric values".to_string(),
            ),
        }
    }
}

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

    fn rem(self, other: BaseValue) -> Self::Output {
        match (self, other) {
            (BaseValue::Number(a), BaseValue::Number(b)) => match a.rem(b) {
                Some(result) => Ok(BaseValue::Number(result)),
                None => runtime_error(SourceSlice::synthetic(), "Division by zero".to_string()),
            },
            _ => runtime_error(
                SourceSlice::synthetic(),
                "Cannot divide non-numeric values".to_string(),
            ),
        }
    }
}

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

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