RustAst/src/ast/environment.rs

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

use crate::ast::compiler::bound_nodes::{Address, BoundNode};
use crate::ast::compiler::dumper::Dumper;
use crate::ast::compiler::lambda_collector::LambdaCollector;
use crate::ast::compiler::macros::{MacroEvaluator, MacroExpander, MacroRegistry};
use crate::ast::compiler::optimizer::{Optimizer, Purity};
use crate::ast::compiler::specializer::{FunctionRegistry, MonoCache, Specializer};
use crate::ast::compiler::tco::{ExecNode, TCO};
use crate::ast::rtl;
use crate::ast::rtl::intrinsics;

pub struct Environment {
    pub global_names: Rc<RefCell<HashMap<Symbol, u32>>>,
    pub global_types: Rc<RefCell<HashMap<u32, StaticType>>>,
    pub global_purity: Rc<RefCell<HashMap<u32, Purity>>>,
    pub global_values: Rc<RefCell<Vec<Value>>>,
    pub prng: Rc<RefCell<fastrand::Rng>>,
    pub function_registry: Rc<RefCell<HashMap<u32, BoundNode>>>,
    pub typed_function_registry: Rc<RefCell<HashMap<u32, TypedNode>>>,
    pub monomorph_cache: Rc<RefCell<MonoCache>>,
    pub debug_mode: bool,
    pub optimization: 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 exec_ast = TCO::optimize(typed_ast);

        let mut vm = VM::new(self.global_values.clone());
        vm.run(&exec_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_purity: Rc::new(RefCell::new(HashMap::new())),
            global_values: Rc::new(RefCell::new(Vec::new())),
            prng: Rc::new(RefCell::new(fastrand::Rng::new())),
            function_registry: Rc::new(RefCell::new(HashMap::new())),
            typed_function_registry: Rc::new(RefCell::new(HashMap::new())),
            monomorph_cache: Rc::new(RefCell::new(HashMap::new())),
            debug_mode: false,
            optimization: true,
        };
        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,
        purity_level: Purity,
        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 mut purity = self.global_purity.borrow_mut();

        let idx = values.len() as u32;
        names.insert(Symbol::from(name), idx);
        types.insert(idx, ty);
        purity.insert(idx, purity_level);
        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 mut purity = self.global_purity.borrow_mut();

        let idx = values.len() as u32;
        names.insert(Symbol::from(name), idx);
        types.insert(idx, ty);
        purity.insert(idx, Purity::Pure); // Constants are always pure
        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, &[])?;

        // 7. Collect Typed Lambdas
        LambdaCollector::collect(&typed_ast, &mut self.typed_function_registry.borrow_mut());

        Ok(typed_ast)
    }

    /// Backend: Optimization (TCO, etc.)
    pub fn link(&self, node: TypedNode) -> ExecNode {
        // 1. Specialize (Always performed for correctness)
        let specialized = self.specialize_node(node);

        // 2. Optimize (Level 1: Cracking, Level 2: Collapsing)
        let optimizer = Optimizer::new(self.optimization)
            .with_globals(self.global_values.clone())
            .with_purity(self.global_purity.clone())
            .with_registry(self.typed_function_registry.clone());
        let optimized = optimizer.optimize(specialized);

        // 3. TCO (Always performed, converts to ExecNode)
        TCO::optimize(optimized)
    }

    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 global_purity = self.global_purity.clone();
        let optimization = self.optimization;

        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 (Phase 2: Cracking & Folding)
                let optimizer = Optimizer::new(optimization)
                    .with_globals(global_values.clone())
                    .with_purity(global_purity.clone());
                let optimized_ast = optimizer.optimize(specialized_ast);

                // 4. TCO (converts to ExecNode)
                let tco_ast = TCO::optimize(optimized_ast);

                // 5. Compile to Value (VM)
                let mut vm = VM::new(global_values.clone());
                let compiled_val = match vm.run(&tco_ast) {
                    Ok(v) => v,
                    Err(e) => return Err(format!("VM Error during specialization: {}", e)),
                };

                // 6. Determine correct return type from the newly inferred function signature
                let ret_type = if let StaticType::Function(sig) = &tco_ast.ty.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: &ExecNode) -> Result<Value, String> {
        let mut vm = VM::new(self.global_values.clone());
        let mut result = vm.run(node)?;

        // Handle potential script body closure
        if let Value::Object(obj) = &result
            && let Some(closure) = obj.as_any().downcast_ref::<crate::ast::vm::Closure>()
        {
            result = vm.run(&closure.exec_node)?;
        }

        // IMPORTANT: Resolve any pending tail call requests from the top-level execution
        while let Value::TailCallRequest(payload) = result {
            let (next_obj, next_args) = *payload;
            if let Some(closure) = next_obj.as_any().downcast_ref::<crate::ast::vm::Closure>() {
                result = vm.run_with_args(closure, next_args)?;
            } else {
                return Err(format!(
                    "Tail call target is not a closure: {}",
                    next_obj.type_name()
                ));
            }
        }
        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
            && let Some(closure) = obj.as_any().downcast_ref::<crate::ast::vm::Closure>()
        {
            result = vm.run_with_observer(&mut observer, &closure.exec_node);
        }

        // Resolve top-level tail calls
        while let Ok(Value::TailCallRequest(payload)) = result {
            let (next_obj, next_args) = *payload;
            if let Some(closure) = next_obj.as_any().downcast_ref::<crate::ast::vm::Closure>() {
                result = vm.run_with_args_observed(&mut observer, closure, next_args);
            } else {
                result = Err(format!(
                    "Tail call target is not a closure: {}",
                    next_obj.type_name()
                ));
                break;
            }
        }

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