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::lambda_collector::LambdaCollector;
use crate::ast::compiler::dumper::Dumper;
use crate::ast::compiler::macros::{MacroExpander, MacroRegistry, MacroEvaluator};
use crate::ast::compiler::specializer::{Specializer, MonoCache, FunctionRegistry};
use crate::ast::rtl;
use crate::ast::rtl::intrinsics;
use crate::ast::compiler::bound_nodes::{Address, BoundNode};

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 function_registry: Rc<RefCell<HashMap<u32, BoundNode>>>,
    pub monomorph_cache: Rc<RefCell<MonoCache>>,
    pub debug_mode: bool,
}

struct EnvFunctionRegistry {
    registry: Rc<RefCell<HashMap<u32, BoundNode>>>,
}

impl FunctionRegistry for EnvFunctionRegistry {
    fn resolve(&self, addr: Address) -> Option<BoundNode> {
        if let Address::Global(idx) = addr {
            self.registry.borrow().get(&idx).cloned()
        } else {
            None
        }
    }
}

/// 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())),
            function_registry: Rc::new(RefCell::new(HashMap::new())),
            monomorph_cache: Rc::new(RefCell::new(HashMap::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. Collect Lambdas (Populate the registry with untyped templates)
        LambdaCollector::collect(&bound_ast, &mut self.function_registry.borrow_mut());

        // 6. 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 {
        // 1. Specialize
        let specialized = self.specialize_node(node);

        // 2. Optimize
        TCO::optimize(specialized)
    }

    fn specialize_node(&self, node: TypedNode) -> TypedNode {
        let registry = Rc::new(EnvFunctionRegistry {
            registry: self.function_registry.clone(),
        });
        
        let rtl_lookup = Rc::new(|name: &str, args: &[StaticType]| intrinsics::lookup(name, args));
        
        let func_reg = self.function_registry.clone();
        let mono_cache = self.monomorph_cache.clone();
        let global_values = self.global_values.clone();
        let global_types = self.global_types.clone();
        
        let compiler = Rc::new(move |func_template: BoundNode, arg_types: &[StaticType]| -> Result<(Value, StaticType), String> {
            // 1. Re-TypeCheck the template with concrete argument types
            let checker = TypeChecker::new(global_types.clone());
            let retyped_ast = checker.check(func_template, arg_types)?; 
            
            // 2. Specialize (Recursive)
            let sub_registry = Rc::new(EnvFunctionRegistry { registry: func_reg.clone() });
            let sub_rtl_lookup = Rc::new(|name: &str, args: &[StaticType]| intrinsics::lookup(name, args));
            
            let sub_specializer = Specializer::new(
                Some(sub_registry),
                None, 
                Some(sub_rtl_lookup),
                Some(mono_cache.clone())
            );
            
            let specialized_ast = sub_specializer.specialize(retyped_ast);
            
            // 3. Optimize (TCO)
            let optimized_ast = TCO::optimize(specialized_ast);
            
            // 4. Compile to Value (VM)
            let mut vm = VM::new(global_values.clone());
            let compiled_val = match vm.run(&optimized_ast) {
                Ok(v) => v,
                Err(e) => return Err(format!("VM Error during specialization: {}", e)),
            };

            // 5. Determine correct return type from the newly inferred function signature
            let ret_type = if let StaticType::Function(sig) = &optimized_ast.ty {
                sig.ret.clone()
            } else {
                StaticType::Any
            };

            Ok((compiled_val, ret_type))
        });
        
        let specializer = Specializer::new(
            Some(registry),
            Some(compiler),
            Some(rtl_lookup),
            Some(self.monomorph_cache.clone()),
        );
        
        specializer.specialize(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());
        let result = vm.run(node)?;
        
        if let Value::Object(obj) = &result {
            if let Some(closure) = obj.as_any().downcast_ref::<crate::ast::vm::Closure>() {
                // Execute the script body
                return vm.run(&closure.function_node);
            }
        }
        
        Ok(result)
    }

    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 mut result = vm.run_with_observer(&mut observer, &linked);
        
        // If result is a closure (script entry), execute the body too
        if let Ok(Value::Object(obj)) = &result {
            if let Some(closure) = obj.as_any().downcast_ref::<crate::ast::vm::Closure>() {
                result = vm.run_with_observer(&mut observer, &closure.function_node);
            }
        }

        Ok((result, observer.logs))
    }
}