RustAst/src/ast/environment.rs

use std::rc::Rc;
use std::cell::RefCell;
use std::collections::HashMap;
use crate::ast::types::{Value, StaticType, Object};
use crate::ast::nodes::{Node, UntypedKind, Symbol};
use crate::ast::parser::Parser;
use crate::ast::compiler::binder::Binder;
use crate::ast::compiler::{TypedNode, TypeChecker};
use crate::ast::vm::{VM, TracingObserver};

use crate::ast::compiler::tco::TCO;
use crate::ast::compiler::dumper::Dumper;
use crate::ast::compiler::macros::{MacroExpander, MacroRegistry, MacroEvaluator};
use crate::ast::rtl;

pub struct Environment {
    pub global_names: Rc<RefCell<HashMap<Symbol, u32>>>,
    pub global_types: Rc<RefCell<HashMap<u32, StaticType>>>,
    pub global_values: Rc<RefCell<Vec<Value>>>,
    pub debug_mode: bool,
}

/// Evaluator used during macro expansion to allow compile-time logic.
struct RuntimeMacroEvaluator {
    global_names: Rc<RefCell<HashMap<Symbol, u32>>>,
    global_types: Rc<RefCell<HashMap<u32, StaticType>>>,
    global_values: Rc<RefCell<Vec<Value>>>,
}

impl MacroEvaluator for RuntimeMacroEvaluator {
    fn evaluate(&self, node: &Node<UntypedKind>, bindings: &HashMap<Rc<str>, Node<UntypedKind>>) -> Result<Value, String> {
        // 1. Check if it's a simple parameter substitution
        if let UntypedKind::Identifier(sym) = &node.kind 
            && let Some(arg_node) = bindings.get(&sym.name) {
                return Ok(Value::Object(Rc::new(arg_node.clone()) as Rc<dyn Object>));
        }

        // 2. Full evaluation for complex compile-time expressions
        let bound_ast = Binder::bind_root(self.global_names.clone(), node)?;
        
        let checker = TypeChecker::new(self.global_types.clone());
        let typed_ast = checker.check(bound_ast)?;
        
        let mut vm = VM::new(self.global_values.clone());
        vm.run(&typed_ast)
    }
}

impl Default for Environment {
    fn default() -> Self {
        Self::new()
    }
}

impl Environment {
    pub fn new() -> Self {
        let env = Self {
            global_names: Rc::new(RefCell::new(HashMap::new())),
            global_types: Rc::new(RefCell::new(HashMap::new())),
            global_values: Rc::new(RefCell::new(Vec::new())),
            debug_mode: false,
        };
        env.register_stdlib();
        env
    }

    pub fn set_debug_mode(&mut self, enabled: bool) {
        self.debug_mode = enabled;
    }

    fn get_expander(&self) -> MacroExpander<RuntimeMacroEvaluator> {
        let evaluator = RuntimeMacroEvaluator {
            global_names: self.global_names.clone(),
            global_types: self.global_types.clone(),
            global_values: self.global_values.clone(),
        };
        MacroExpander::new(MacroRegistry::new(), evaluator)
    }

    pub fn register_native(&self, name: &str, ty: StaticType, func: impl Fn(Vec<Value>) -> Value + 'static) {
        let mut names = self.global_names.borrow_mut();
        let mut types = self.global_types.borrow_mut();
        let mut values = self.global_values.borrow_mut();

        let idx = values.len() as u32;
        names.insert(Symbol::from(name), idx);
        types.insert(idx, ty);
        values.push(Value::Function(Rc::new(func)));
    }

    pub fn register_constant(&self, name: &str, ty: StaticType, val: Value) {
        let mut names = self.global_names.borrow_mut();
        let mut types = self.global_types.borrow_mut();
        let mut values = self.global_values.borrow_mut();

        let idx = values.len() as u32;
        names.insert(Symbol::from(name), idx);
        types.insert(idx, ty);
        values.push(val);
    }

    fn register_stdlib(&self) {
        // Register all standard library functions via RTL module
        rtl::register(self);
    }

    pub fn dump_ast(&self, source: &str) -> Result<String, String> {
        let compiled = self.compile(source)?;
        let linked = self.link(compiled);
        Ok(Dumper::dump(&linked))
    }

    /// Frontend: Parse -> Expand Macros -> Bind -> Type Check
    pub fn compile(&self, source: &str) -> Result<TypedNode, String> {
        // 1. Parse
        let mut parser = Parser::new(source)?;
        let untyped_ast = parser.parse_expression()?;

        // 2. Check for trailing tokens
        if !parser.at_eof() {
            return Err("Unexpected trailing expressions in script. Use (do ...) for sequences.".to_string());
        }

        // 3. Expand Macros
        let expanded_ast = self.get_expander().expand(untyped_ast)?;

        // 4. Bind
        let bound_ast = Binder::bind_root(self.global_names.clone(), &expanded_ast)?;

        // 5. Type Check
        let checker = TypeChecker::new(self.global_types.clone());
        let typed_ast = checker.check(bound_ast)?;

        Ok(typed_ast)
    }

    /// Backend: Optimization (TCO, etc.)
    pub fn link(&self, node: TypedNode) -> TypedNode {
        TCO::optimize(node)
    }

    /// Runtime: Execute the linked AST in the VM
    pub fn run(&self, node: &TypedNode) -> Result<Value, String> {
        let mut vm = VM::new(self.global_values.clone());
        vm.run(node)
    }

    pub fn run_script(&self, source: &str) -> Result<Value, String> {
        if self.debug_mode {
            let (res, logs) = self.run_debug(source)?;
            for line in logs {
                println!("{}", line);
            }
            res
        } else {
            let compiled = self.compile(source)?;
            let linked = self.link(compiled);
            self.run(&linked)
        }
    }

    pub fn run_debug(&self, source: &str) -> Result<(Result<Value, String>, Vec<String>), String> {
        let compiled = self.compile(source)?;
        let linked = self.link(compiled);

        // Execute with TracingObserver
        let mut vm = VM::new(self.global_values.clone());
        let mut observer = TracingObserver::new();
        let result = vm.run_with_observer(&mut observer, &linked);
        
        Ok((result, observer.logs))
    }
}