RustAst/src/ast/types.rs

use std::collections::HashMap;
use std::rc::Rc;
use std::cell::RefCell;
use std::fmt;
use std::sync::OnceLock;
use std::sync::Mutex; // Still needed for global keyword registry
use std::any::Any;
use chrono::{TimeZone, Utc};

/// Simple source location
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct SourceLocation {
    pub line: u32,
    pub col: u32,
}

/// Shared identity for nodes (Location, etc.)
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct NodeIdentity {
    pub location: SourceLocation,
}

pub type Identity = Rc<NodeIdentity>;

/// Interned string identifier
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
pub struct Keyword(pub u32);

static KEYWORD_REGISTRY: OnceLock<Mutex<HashMap<String, u32>>> = OnceLock::new();
static KEYWORD_REVERSE: OnceLock<Mutex<Vec<String>>> = OnceLock::new();

impl Keyword {
    pub fn intern(name: &str) -> Self {
        let mut reg = KEYWORD_REGISTRY.get_or_init(|| Mutex::new(HashMap::new())).lock().unwrap();
        if let Some(&id) = reg.get(name) {
            Keyword(id)
        } else {
            let mut rev = KEYWORD_REVERSE.get_or_init(|| Mutex::new(Vec::new())).lock().unwrap();
            let id = rev.len() as u32;
            reg.insert(name.to_string(), id);
            rev.push(name.to_string());
            Keyword(id)
        }
    }

    pub fn name(&self) -> String {
        let rev = KEYWORD_REVERSE.get_or_init(|| Mutex::new(Vec::new())).lock().unwrap();
        rev[self.0 as usize].clone()
    }
}

/// Interface for custom objects (Closures, Series, Streams)
pub trait Object: fmt::Debug {
    fn type_name(&self) -> &'static str;
    fn as_any(&self) -> &dyn Any;
}

/// Internal storage for Tuples (and Records) to allow sharing schema (keys) between instances.
#[derive(Debug, Clone)]
pub struct TupleData {
    pub values: Vec<Value>,
    /// Optional names for slots. If present, this is semantically a Record.
    pub keys: Option<Rc<Vec<Keyword>>>,
}

/// Core data value in Myc Script (similar to TDataValue)
#[derive(Clone)]
pub enum Value {
    Void,
    Bool(bool),
    Int(i64),
    Float(f64),
    DateTime(i64),
    Text(Rc<str>),
    Keyword(Keyword),
    Tuple(Rc<TupleData>), // Replaces List and Record
    Function(Rc<dyn Fn(Vec<Value>) -> Value>),
    Object(Rc<dyn Object>), // For compiled Closures and other opaque types
    Cell(Rc<RefCell<Value>>), // Boxed value for captures
    TailCallRequest(Box<(Rc<dyn Object>, Vec<Value>)>), // Internal: For TCO (Boxed to keep Value small)
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub struct Signature {
    pub params: StaticType,
    pub ret: StaticType,
}

#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum StaticType {
    Any,
    Void,
    Bool,
    Int,
    Float,
    DateTime,
    Text,
    Keyword,
    List(Box<StaticType>), // Legacy / Dynamic list
    Tuple(Vec<StaticType>), // Heterogeneous fixed-size
    Vector(Box<StaticType>, usize), // Homogeneous fixed-size
    Matrix(Box<StaticType>, Vec<usize>), // Multi-dimensional homogeneous
    Record(Rc<Vec<(Keyword, StaticType)>>),
    Function(Box<Signature>),
    FunctionOverloads(Vec<Signature>),
    Object(&'static str),
}

impl fmt::Display for StaticType {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            StaticType::Any => write!(f, "any"),
            StaticType::Void => write!(f, "void"),
            StaticType::Bool => write!(f, "bool"),
            StaticType::Int => write!(f, "int"),
            StaticType::Float => write!(f, "float"),
            StaticType::DateTime => write!(f, "datetime"),
            StaticType::Text => write!(f, "text"),
            StaticType::Keyword => write!(f, "keyword"),
            StaticType::List(inner) => write!(f, "[{}]", inner),
            StaticType::Tuple(elements) => {
                write!(f, "[")?;
                for (i, el) in elements.iter().enumerate() {
                    if i > 0 { write!(f, " ")?; }
                    write!(f, "{}", el)?;
                }
                write!(f, "]")
            },
            StaticType::Vector(inner, len) => write!(f, "vector<{}, {}>", inner, len),
            StaticType::Matrix(inner, shape) => {
                write!(f, "matrix<{}, [", inner)?;
                for (i, s) in shape.iter().enumerate() {
                    if i > 0 { write!(f, " ")?; }
                    write!(f, "{}", s)?;
                }
                write!(f, "]>")
            },
            StaticType::Record(fields) => {
                write!(f, "{{")?;
                for (i, (k, v)) in fields.iter().enumerate() {
                    if i > 0 { write!(f, ", ")?; }
                    write!(f, ":{} {}", k.name(), v)?;
                }
                write!(f, "}}")
            },
            StaticType::Function(sig) => {
                write!(f, "fn({}) -> {}", sig.params, sig.ret)
            },
            StaticType::FunctionOverloads(sigs) => {
                write!(f, "overloads({} variants)", sigs.len())
            }
            StaticType::Object(name) => write!(f, "{}", name),
        }
    }
}

impl StaticType {
    /// Returns true if `other` can be assigned to a location of type `self`.
    pub fn is_assignable_from(&self, other: &StaticType) -> bool {
        if self == other || matches!(self, StaticType::Any) || matches!(other, StaticType::Any) {
            return true;
        }

        match (self, other) {
            // A Vector is a Tuple
            (StaticType::Tuple(elements), StaticType::Vector(inner, len)) => {
                if elements.len() != *len { return false; }
                elements.iter().all(|e| e.is_assignable_from(inner))
            },
            // A Matrix is a Vector (of Vectors/Matrices)
            (StaticType::Vector(inner, len), StaticType::Matrix(m_inner, shape)) => {
                if shape.is_empty() || shape[0] != *len { return false; }
                if shape.len() == 1 {
                    inner.is_assignable_from(m_inner)
                } else {
                    // It's a matrix of higher dimension, so inner must be assignable from a sub-matrix
                    let sub_shape = shape[1..].to_vec();
                    inner.is_assignable_from(&StaticType::Matrix(m_inner.clone(), sub_shape))
                }
            },
            _ => false
        }
    }

    /// Tries to resolve a call with the given argument type (usually a Tuple) and returns the return type.
    pub fn resolve_call(&self, args_ty: &StaticType) -> Option<StaticType> {
        match self {
            StaticType::Any => Some(StaticType::Any),
            StaticType::Function(sig) => {
                if sig.params.is_assignable_from(args_ty) {
                    Some(sig.ret.clone())
                } else {
                    None
                }
            }
            StaticType::FunctionOverloads(sigs) => {
                sigs.iter()
                    .find(|sig| sig.params.is_assignable_from(args_ty))
                    .map(|sig| sig.ret.clone())
            }
            _ => None,
        }
    }

    /// Returns true if this type and all its recursive elements are scalars (Int, Float, etc.)
    pub fn is_scalar_pure(&self) -> bool {
        match self {
            StaticType::Int | StaticType::Float | StaticType::Bool | StaticType::DateTime => true,
            StaticType::Tuple(elements) => elements.iter().all(|e| e.is_scalar_pure()),
            StaticType::Vector(inner, _) => inner.is_scalar_pure(),
            StaticType::Matrix(inner, _) => inner.is_scalar_pure(),
            _ => false,
        }
    }
}

impl Value {
    pub fn is_truthy(&self) -> bool {
        match self {
            Value::Void => false,
            Value::Bool(b) => *b,
            Value::Cell(c) => c.borrow().is_truthy(),
            _ => true,
        }
    }

    pub fn make_list(values: Vec<Value>) -> Self {
        Value::Tuple(Rc::new(TupleData { values, keys: None }))
    }

    pub fn make_record(keys: Vec<Keyword>, values: Vec<Value>) -> Self {
        Value::Tuple(Rc::new(TupleData { values, keys: Some(Rc::new(keys)) }))
    }

    pub fn static_type(&self) -> StaticType {
        match self {
            Value::Void => StaticType::Void,
            Value::Bool(_) => StaticType::Bool,
            Value::Int(_) => StaticType::Int,
            Value::Float(_) => StaticType::Float,
            Value::DateTime(_) => StaticType::DateTime,
            Value::Text(_) => StaticType::Text,
            Value::Keyword(_) => StaticType::Keyword,
            Value::Tuple(t) => {
                if let Some(keys) = &t.keys {
                    // It's a Record
                    let mut fields = Vec::with_capacity(t.values.len());
                    for (i, v) in t.values.iter().enumerate() {
                        fields.push((keys[i], v.static_type()));
                    }
                    StaticType::Record(Rc::new(fields))
                } else {
                    // It's a List/Tuple
                    let l = &t.values;
                    if l.is_empty() {
                        return StaticType::Vector(Box::new(StaticType::Any), 0);
                    }
                    
                    let element_types: Vec<_> = l.iter().map(|v| v.static_type()).collect();
                    
                    // Check for Homogeneity (Vector)
                    let first_ty = &element_types[0];
                    let all_same = element_types.iter().all(|t| t == first_ty);
                    
                    if all_same {
                        match first_ty {
                            StaticType::Vector(inner, len) => {
                                // Possible Matrix
                                StaticType::Matrix(inner.clone(), vec![l.len(), *len])
                            },
                            StaticType::Matrix(inner, shape) => {
                                let mut new_shape = vec![l.len()];
                                new_shape.extend(shape);
                                StaticType::Matrix(inner.clone(), new_shape)
                            },
                            _ => StaticType::Vector(Box::new(first_ty.clone()), l.len())
                        }
                    } else {
                        StaticType::Tuple(element_types)
                    }
                }
            },
            Value::Function(_) => StaticType::Any, // Dynamic function
            Value::Object(o) => StaticType::Object(o.type_name()),
            Value::Cell(c) => c.borrow().static_type(),
            Value::TailCallRequest(_) => StaticType::Any, // Internal state, but typable as Any
        }
    }
}

impl fmt::Display for Value {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Value::Void => write!(f, "void"),
            Value::Bool(b) => write!(f, "{}", b),
            Value::Int(i) => write!(f, "{}", i),
            Value::Float(fl) => write!(f, "{}", fl),
            Value::DateTime(ts) => {
                match Utc.timestamp_millis_opt(*ts) {
                    chrono::LocalResult::Single(dt) => write!(f, "#{}#", dt.format("%Y-%m-%d %H:%M:%S")),
                    _ => write!(f, "#timestamp({})#", ts),
                }
            },
            Value::Text(t) => write!(f, "\"{}\"", t),
            Value::Keyword(k) => write!(f, ":{}", k.name()),
            Value::Tuple(t) => {
                if let Some(keys) = &t.keys {
                    write!(f, "{{")?;
                    for i in 0..t.values.len() {
                        if i > 0 { write!(f, ", ")?; }
                        write!(f, ":{} {}", keys[i].name(), t.values[i])?;
                    }
                    write!(f, "}}")
                } else {
                    write!(f, "[")?;
                    for (i, val) in t.values.iter().enumerate() {
                        if i > 0 { write!(f, " ")?; }
                        write!(f, "{}", val)?;
                    }
                    write!(f, "]")
                }
            },
            Value::Function(_) => write!(f, "<native fn>"),
            Value::Object(o) => write!(f, "<{}>", o.type_name()),
            Value::Cell(c) => write!(f, "{}", c.borrow()),
            Value::TailCallRequest(_) => write!(f, "<tail call request>"),
        }
    }
}

impl fmt::Debug for Value {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        fmt::Display::fmt(self, f)
    }
}