RustAst/src/ast/environment.rs

use crate::ast::compiler::analyzer::Analyzer;
use crate::ast::compiler::binder::Binder;
use crate::ast::compiler::{TypeChecker, TypedNode};
use crate::ast::nodes::{Symbol, SyntaxKind, SyntaxNode};
use crate::ast::parser::Parser;
use crate::ast::vm::{TracingObserver, VM};
use std::cell::RefCell;
use std::collections::{HashMap, HashSet};
use std::path::{Path, PathBuf};
use std::rc::Rc;

use crate::ast::compiler::bound_nodes::{
    Address, AnalyzedNode, BoundKind, GlobalAnalyzedRegistry, GlobalFunctionRegistry, GlobalIdx,
    Node, 
};
use crate::ast::compiler::dumper::Dumper;
use crate::ast::compiler::lambda_collector::LambdaCollector;
use crate::ast::compiler::lowering::{ExecNode, Lowering};
use crate::ast::compiler::macros::{MacroEvaluator, MacroExpander, MacroRegistry};
use crate::ast::compiler::optimizer::Optimizer;
use crate::ast::compiler::specializer::{FunctionRegistry, MonoCache, Specializer};
use crate::ast::rtl::intrinsics;
use crate::ast::types::{
    NodeIdentity, Object, Purity, SourceLocation, StaticType, Value,
};

use crate::ast::diagnostics::{DiagnosticLevel, Diagnostics};
pub type PipelineGenerator = Box<dyn FnMut() -> bool>;

const SYSTEM_LIB_SOURCE: &str = include_str!("system.myc");
const SYSTEM_LIB_PATH: &str = "<embedded>/system.myc";

pub struct CompilationResult {
    pub ast: Option<TypedNode>,
    pub diagnostics: Diagnostics,
}

impl CompilationResult {
    pub fn success(ast: TypedNode) -> Self {
        Self {
            ast: Some(ast),
            diagnostics: Diagnostics::new(),
        }
    }

    pub fn error(msg: impl Into<String>) -> Self {
        let mut diag = Diagnostics::new();
        diag.push_error(msg, None);
        Self {
            ast: None,
            diagnostics: diag,
        }
    }

    pub fn into_result(self) -> Result<TypedNode, String> {
        if self.diagnostics.has_errors() {
            let errors: Vec<String> = self
                .diagnostics
                .items
                .into_iter()
                .filter(|d| d.level == DiagnosticLevel::Error)
                .map(|d| d.message)
                .collect();
            Err(errors.join("\n"))
        } else if let Some(ast) = self.ast {
            Ok(ast)
        } else {
            Err("Compilation failed without diagnostics".to_string())
        }
    }
}

pub struct Environment {
    pub root_types: Rc<RefCell<Vec<StaticType>>>,
    pub root_purity: Rc<RefCell<Vec<Purity>>>,
    pub root_values: Rc<RefCell<Vec<Value>>>,
    pub fixed_scope_idx: i32,
    pub root_scopes: Rc<RefCell<Vec<crate::ast::compiler::binder::CompilerScope>>>,
    pub root_slot_count: Rc<RefCell<u32>>,
    pub function_registry: Rc<RefCell<GlobalFunctionRegistry>>,
    pub typed_function_registry: Rc<RefCell<GlobalAnalyzedRegistry>>,
    pub monomorph_cache: Rc<RefCell<MonoCache>>,
    pub debug_mode: bool,
    pub optimization: bool,
    pub macro_registry: Rc<RefCell<MacroRegistry>>,
    pub pipeline_generators: Rc<RefCell<Vec<PipelineGenerator>>>,
    pub search_paths: Rc<RefCell<Vec<PathBuf>>>,
    pub loaded_modules: Rc<RefCell<HashSet<PathBuf>>>,
}

struct EnvFunctionRegistry {
    registry: Rc<RefCell<GlobalFunctionRegistry>>,
    analyzed_registry: Rc<RefCell<GlobalAnalyzedRegistry>>,
}

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

struct RuntimeMacroEvaluator {
    root_scopes: Rc<RefCell<Vec<crate::ast::compiler::binder::CompilerScope>>>,
    root_slot_count: Rc<RefCell<u32>>,
    root_types: Rc<RefCell<Vec<StaticType>>>,
    root_values: Rc<RefCell<Vec<Value>>>,
    root_purity: Rc<RefCell<Vec<Purity>>>,
    fixed_scope_idx: i32,
}

impl MacroEvaluator for RuntimeMacroEvaluator {
    fn evaluate(
        &self,
        node: &SyntaxNode,
        bindings: &HashMap<Rc<str>, SyntaxNode>,
    ) -> Result<Value, String> {
        if let SyntaxKind::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>));
        }

        let mut diag = Diagnostics::new();
        let initial_scopes = self.root_scopes.borrow().clone();
        let initial_slot_count = *self.root_slot_count.borrow();

        let (bound_ast, captures, _, _) = Binder::bind_root(initial_scopes, initial_slot_count, self.fixed_scope_idx, node, &mut diag)?;
        let bound_ast = crate::ast::compiler::captures::CapturePass::apply(bound_ast, &captures);
        let checker = TypeChecker::new(self.root_types.clone());
        let typed_ast = checker.check(&bound_ast, &[], &mut diag);

        if diag.has_errors() {
            return Err(diag
                .items
                .into_iter()
                .map(|d| d.message)
                .collect::<Vec<_>>()
                .join("\n"));
        }

        let exec_ast = Lowering::lower(Analyzer::analyze(&typed_ast, &self.root_purity.borrow())); // Minimal analysis for macro eval

        let mut vm = VM::new(self.root_values.clone());
        vm.run(&exec_ast)
    }
}

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

impl Environment {
    pub fn new() -> Self {
        let env = Self {
            root_types: Rc::new(RefCell::new(Vec::new())),
            root_purity: Rc::new(RefCell::new(Vec::new())),
            root_values: Rc::new(RefCell::new(Vec::new())),
            fixed_scope_idx: -1,
            root_scopes: Rc::new(RefCell::new(vec![crate::ast::compiler::binder::CompilerScope::new()])),
            root_slot_count: Rc::new(RefCell::new(0)),
            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,
            macro_registry: Rc::new(RefCell::new(MacroRegistry::new())),
            pipeline_generators: Rc::new(RefCell::new(Vec::new())),
            search_paths: Rc::new(RefCell::new(Vec::new())),
            loaded_modules: Rc::new(RefCell::new(HashSet::new())),
        };
        crate::ast::rtl::register(&env);
        
        let mut env = env;
        env.fixed_scope_idx = 0;
        // Push the first mutable user scope (Level 1)
        env.root_scopes.borrow_mut().push(crate::ast::compiler::binder::CompilerScope::new());

        // Automatically add standard search paths (CWD and CWD/rtl)
        if let Ok(cwd) = std::env::current_dir() {
            env.add_search_path(&cwd);
            let rtl_path = cwd.join("rtl");
            if rtl_path.exists() {
                env.add_search_path(rtl_path);
            }
        }

        env
    }

    pub fn add_search_path(&self, path: impl AsRef<Path>) {
        self.search_paths.borrow_mut().push(path.as_ref().to_path_buf());
    }

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

    /// Pumps data through all registered pipeline generators until they are exhausted.
    pub fn run_pipeline(&self) {
        let mut generators = self.pipeline_generators.borrow_mut();
        let mut any_active = true;
        while any_active {
            any_active = false;
            for generator in generators.iter_mut() {
                if generator() {
                    any_active = true;
                }
            }
        }
    }

    fn get_expander(&self) -> MacroExpander<RuntimeMacroEvaluator> {
        let evaluator = RuntimeMacroEvaluator {
            root_scopes: self.root_scopes.clone(),
            root_slot_count: self.root_slot_count.clone(),
            root_types: self.root_types.clone(),
            root_values: self.root_values.clone(),
            root_purity: self.root_purity.clone(),
            fixed_scope_idx: self.fixed_scope_idx,
        };
        MacroExpander::new(self.macro_registry.borrow().clone(), evaluator)
    }


    /// Resolves a #use module path relative to a base path, then falls back to search paths.
    /// Returns a list of all matching .myc files (single file, or all files in a directory).
    fn resolve_module_paths(&self, module_path: &str, base_path: &Path) -> Result<Vec<PathBuf>, String> {
        let relative_path = module_path.replace("->", std::path::MAIN_SEPARATOR_STR);
        let file_name = format!("{}.myc", relative_path);

        let check_path = |base: &Path| -> Option<Vec<PathBuf>> {
            // Check for file first
            let mut file_p = base.to_path_buf();
            file_p.push(&file_name);
            if file_p.is_file() && let Ok(canon) = file_p.canonicalize() {
                return Some(vec![canon]);
            }

            // Check for directory
            let mut dir_p = base.to_path_buf();
            dir_p.push(&relative_path);
            if dir_p.is_dir() {
                let mut files = Vec::new();
                if let Ok(entries) = std::fs::read_dir(&dir_p) {
                    let mut valid_entries: Vec<_> = entries.filter_map(|e| e.ok()).collect();
                    valid_entries.sort_by_key(|e| e.path());
                    for entry in valid_entries {
                        let p = entry.path();
                        if p.is_file() && p.extension().is_some_and(|ext| ext == "myc") && let Ok(canon) = p.canonicalize() {
                            files.push(canon);
                        }
                    }
                }
                if !files.is_empty() {
                    return Some(files);
                }
            }
            None
        };

        // 1. Try relative to base_path
        if let Some(paths) = check_path(base_path) {
            return Ok(paths);
        }

        // 2. Try search paths
        for sp in self.search_paths.borrow().iter() {
            if let Some(paths) = check_path(sp) {
                return Ok(paths);
            }
        }

        Err(format!("Could not find module or directory '{}'", module_path))
    }

    fn collect_dependencies(
        &self,
        source: &str,
        base_path: &Path,
        all_files: &mut Vec<(PathBuf, SyntaxNode)>,
    ) -> Result<(), String> {
        let directives = self.extract_use_directives(source);

        for module_path in directives {
            let abs_paths = self.resolve_module_paths(&module_path, base_path)?;

            for abs_path in abs_paths {
                if self.loaded_modules.borrow().contains(&abs_path) {
                    continue;
                }

                self.loaded_modules.borrow_mut().insert(abs_path.clone());

                let lib_source = std::fs::read_to_string(&abs_path)
                    .map_err(|e| format!("Failed to read module {:?}: {}", abs_path, e))?;

                // Pre-process dependencies of the library BEFORE adding it, ensuring topological order
                let lib_base = abs_path.parent().unwrap_or(Path::new("."));
                self.collect_dependencies(&lib_source, lib_base, all_files)?;

                let mut parser = Parser::new(&lib_source);
                let syntax_ast = parser.parse_expression();

                if parser.diagnostics.has_errors() {
                    return Err(format!(
                        "Parser error in module {:?}:\n{}",
                        abs_path,
                        parser.diagnostics.items[0].message
                    ));
                }

                all_files.push((abs_path, syntax_ast));
            }
        }
        Ok(())
    }

    fn extract_use_directives(&self, source: &str) -> Vec<String> {
        let mut paths = Vec::new();
        for line in source.lines() {
            let trimmed = line.trim();
            if trimmed.is_empty() || trimmed.starts_with(';') {
                continue;
            }
            if let Some(stripped) = trimmed.strip_prefix("#use ") {
                let path = stripped.trim();
                // Strip optional quotes if they somehow got in, though user spec says no spaces/quotes.
                let clean_path = if (path.starts_with('"') && path.ends_with('"'))
                    || (path.starts_with('\'') && path.ends_with('\''))
                {
                    &path[1..path.len() - 1]
                } else {
                    path
                };
                paths.push(clean_path.to_string());
            } else {
                // Stop at first non-directive/non-comment line
                break;
            }
        }
        paths
    }

    /// Used to pre-load all dependencies of a script before compiling it.
    /// It returns the base path which can be used to resolve further things if necessary.
    pub fn preload_dependencies(&self, source: &str, file_path: Option<&Path>) -> Result<(), String> {
        let base_path = file_path.and_then(|p| p.parent()).unwrap_or_else(|| Path::new("."));
        let mut files: Vec<(PathBuf, SyntaxNode)> = Vec::new();

        // 1. Always load the embedded system library first as a virtual module
        let system_path = PathBuf::from(SYSTEM_LIB_PATH);
        if !self.loaded_modules.borrow().contains(&system_path) {
            self.loaded_modules.borrow_mut().insert(system_path.clone());
            let mut parser = Parser::new(SYSTEM_LIB_SOURCE);
            let syntax_ast = parser.parse_expression();
            if parser.diagnostics.has_errors() {
                return Err(format!(
                    "Failed to parse embedded system library:\n{}",
                    parser.diagnostics.items[0].message
                ));
            }
            files.push((system_path, syntax_ast));
        }

        // 2. Collect dependencies from the main source
        self.collect_dependencies(source, base_path, &mut files)?;

        // Pass 1: Discovery (Globals and Macros)
        for (_, syntax_ast) in &files {
            self.discover_globals(syntax_ast);
        }

        // Pass 2: Compilation and Initialization
        for (path, syntax_ast) in files {
            let typed_ast = self.compile_syntax(syntax_ast).map_err(|e: String| {
                format!("Compilation error in {}:\n{}", path.display(), e)
            })?;
            self.run_script_compiled(typed_ast).map_err(|e: String| {
                format!("Initialization error in {}:\n{}", path.display(), e)
            })?;
        }

        Ok(())
    }

    fn discover_globals(&self, node: &SyntaxNode) {
        match &node.kind {
            SyntaxKind::Def { target, .. } => {
                if let SyntaxKind::Identifier(sym) = &target.kind {
                    let mut root_scopes = self.root_scopes.borrow_mut();
                    let last_idx = root_scopes.len() - 1;
                    let current_scope = &mut root_scopes[last_idx];

                    if !current_scope.locals.contains_key(sym) {
                        let mut slot_count = self.root_slot_count.borrow_mut();
                        let slot = crate::ast::compiler::bound_nodes::LocalSlot(*slot_count);
                        current_scope.locals.insert(
                            sym.clone(),
                            crate::ast::compiler::binder::LocalInfo {
                                addr: Address::Local(slot),
                                identity: node.identity.clone(),
                                _ty: StaticType::Any,
                                purity: Purity::Impure,
                            },
                        );
                        *slot_count += 1;
                    }
                }
            }
            SyntaxKind::MacroDecl { name, params, body } => {
                let mut registry = self.macro_registry.borrow_mut();

                fn extract_names(node: &SyntaxNode) -> Vec<Rc<str>> {
                    match &node.kind {
                        SyntaxKind::Identifier(sym) => vec![sym.name.clone()],
                        SyntaxKind::Tuple { elements } => {
                            elements.iter().flat_map(extract_names).collect()
                        }
                        _ => vec![],
                    }
                }
                
                let p_names = extract_names(params);
                registry.define(name.name.clone(), p_names, body.as_ref().clone());
            }
            SyntaxKind::Block { exprs } => {
                for expr in exprs {
                    self.discover_globals(expr);
                }
            }
            _ => {}
        }
    }

    fn compile_pipeline(&self, syntax_ast: SyntaxNode, diagnostics: &mut Diagnostics) -> Option<TypedNode> {
        let expanded_ast = match self.get_expander().expand(syntax_ast) {
            Ok(ast) => ast,
            Err(e) => {
                diagnostics.push_error(e, None);
                return None;
            }
        };

        let initial_scopes = self.root_scopes.borrow().clone();
        let initial_slot_count = *self.root_slot_count.borrow();

        let (bound_ast, captures, final_scopes, final_slot_count) =
            match Binder::bind_root(initial_scopes, initial_slot_count, self.fixed_scope_idx, &expanded_ast, diagnostics) {
                Ok(res) => res,
                Err(e) => {
                    diagnostics.push_error(e, None);
                    return None;
                }
            };

        // Update environment state with new bindings from this script
        *self.root_scopes.borrow_mut() = final_scopes;
        *self.root_slot_count.borrow_mut() = final_slot_count;

        let bound_ast = crate::ast::compiler::captures::CapturePass::apply(bound_ast, &captures);

        // Pre-allocate global slots to prevent out-of-bounds during specialization/optimization
        {
            let mut values = self.root_values.borrow_mut();
            let count = *self.root_slot_count.borrow();
            if (count as usize) > values.len() {
                values.resize(count as usize, Value::Void);
            }
        }

        LambdaCollector::collect(&bound_ast, &mut self.function_registry.borrow_mut());

        let checker = TypeChecker::new(self.root_types.clone());
        let wrapped_ast = if let crate::ast::compiler::bound_nodes::BoundKind::Lambda { .. } = bound_ast.kind {
            bound_ast
        } else {
            crate::ast::compiler::bound_nodes::Node {
                identity: bound_ast.identity.clone(),
                kind: crate::ast::compiler::bound_nodes::BoundKind::Lambda {
                    params: std::rc::Rc::new(crate::ast::compiler::bound_nodes::Node {
                        identity: bound_ast.identity.clone(),
                        kind: crate::ast::compiler::bound_nodes::BoundKind::Tuple { elements: vec![] },
                        ty: (),
                    }),
                    upvalues: vec![],
                    body: std::rc::Rc::new(bound_ast),
                    positional_count: Some(0),
                },
                ty: (),
            }
        };
        Some(checker.check(&wrapped_ast, &[], diagnostics))
    }

    fn compile_syntax(&self, syntax_ast: SyntaxNode) -> Result<TypedNode, String> {
        let mut diagnostics = Diagnostics::new();

        let typed_ast_opt = self.compile_pipeline(syntax_ast, &mut diagnostics);

        if diagnostics.has_errors() || typed_ast_opt.is_none() {
            return Err(diagnostics
                .items
                .into_iter()
                .map(|d| d.message)
                .collect::<Vec<_>>()
                .join("\n"));
        }

        Ok(typed_ast_opt.unwrap())
    }

    pub fn register_native(
        &self,
        name: &str,
        ty: StaticType,
        func: Rc<crate::ast::types::NativeFunction>,
    ) {
        let mut types = self.root_types.borrow_mut();
        let mut values = self.root_values.borrow_mut();
        let mut purity = self.root_purity.borrow_mut();

        let identity = NodeIdentity::new(SourceLocation { line: 0, col: 0 });

        // populate root_scopes[0]
        let mut root_scopes = self.root_scopes.borrow_mut();
        let mut slot_count = self.root_slot_count.borrow_mut();
        
        let global_idx = GlobalIdx(*slot_count);
        root_scopes[0].locals.insert(
            Symbol::from(name),
            crate::ast::compiler::binder::LocalInfo {
                addr: Address::Global(global_idx),
                identity,
                _ty: ty.clone(),
                purity: func.purity,
            },
        );
        *slot_count += 1;

        types.push(ty);
        purity.push(func.purity);
        values.push(Value::Function(func));
    }

    pub fn register_native_fn(
        &self,
        name: &str,
        ty: StaticType,
        purity_level: Purity,
        func: impl Fn(&[Value]) -> Value + 'static,
    ) {
        self.register_native(
            name,
            ty,
            Rc::new(crate::ast::types::NativeFunction {
                func: Rc::new(func),
                purity: purity_level,
            }),
        );
    }

    pub fn register_constant(&self, name: &str, ty: StaticType, val: Value) {
        let mut types = self.root_types.borrow_mut();
        let mut values = self.root_values.borrow_mut();
        let mut purity = self.root_purity.borrow_mut();

        let identity = NodeIdentity::new(SourceLocation { line: 0, col: 0 });

        // populate root_scopes[0]
        let mut root_scopes = self.root_scopes.borrow_mut();
        let mut slot_count = self.root_slot_count.borrow_mut();
        
        let global_idx = GlobalIdx(*slot_count);
        root_scopes[0].locals.insert(
            Symbol::from(name),
            crate::ast::compiler::binder::LocalInfo {
                addr: Address::Global(global_idx),
                identity,
                _ty: ty.clone(),
                purity: Purity::Pure,
            },
        );
        *slot_count += 1;

        types.push(ty);
        purity.push(Purity::Pure);
        values.push(val);
    }


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

    pub fn compile(&self, source: &str) -> CompilationResult {
        if let Err(e) = self.preload_dependencies(source, None) {
            return CompilationResult::error(format!("Dependency Error: {}", e));
        }

        let mut parser = Parser::new(source);
        let syntax_ast = parser.parse_expression();

        if !parser.at_eof() {
            parser
                .diagnostics
                .push_error("Unexpected trailing expressions in script.", None);
            return CompilationResult {
                ast: None,
                diagnostics: parser.diagnostics,
            };
        }
        let mut diagnostics = parser.diagnostics;

        let typed_ast = self.compile_pipeline(syntax_ast, &mut diagnostics);

        CompilationResult {
            ast: typed_ast,
            diagnostics,
        }
    }

    pub fn link(&self, node: TypedNode) -> ExecNode {
        // 1. Analyze
        let analyzed = Analyzer::analyze(&node, &self.root_purity.borrow());

        // 2. Collect Analyzed Lambdas
        LambdaCollector::collect(&analyzed, &mut self.typed_function_registry.borrow_mut());

        // 3. Specialize
        let specialized = self.specialize_node(analyzed);

        // 4. Optimize
        let optimizer = Optimizer::new(self.optimization)
            .with_globals(self.root_values.clone())
            .with_purity(self.root_purity.clone())
            .with_registry(self.typed_function_registry.clone());
        let optimized = optimizer.optimize(specialized);

        // 5. Lowering
        Lowering::lower(optimized)
    }

    pub fn instantiate(&self, node: ExecNode) -> Rc<crate::ast::types::NativeFunction> {
        let root_values = self.root_values.clone();
        if let BoundKind::Lambda {
            params,
            upvalues,
            body,
            positional_count,
        } = &node.kind
            && upvalues.is_empty()
        {
            let closure = Rc::new(crate::ast::vm::Closure::new(
                params.clone(),
                body.ty.original.clone(),
                body.clone(),
                Vec::new(),
                *positional_count,
                node.ty.stack_size,
            ));
            let closure_obj: Rc<dyn crate::ast::types::Object> = closure;

            return Rc::new(crate::ast::types::NativeFunction {
                purity: Purity::Impure,
                func: Rc::new(move |args| {
                    let mut vm = VM::new(root_values.clone());
                    match vm.run_with_args(closure_obj.clone(), args) {
                        Ok(v) => v,
                        Err(e) => panic!("Myc Runtime Error: {}", e),
                    }
                }),
            });
        }

        let exec_node = Rc::new(node);
        Rc::new(crate::ast::types::NativeFunction {
            purity: Purity::Impure,
            func: Rc::new(move |args| {
                let mut vm = VM::new(root_values.clone());
                let res = match vm.run(&exec_node) {
                    Ok(v) => v,
                    Err(e) => panic!("Myc Runtime Error: {}", e),
                };

                if let Value::Object(obj) = &res
                    && obj
                        .as_any()
                        .downcast_ref::<crate::ast::vm::Closure>()
                        .is_some()
                {
                    match vm.run_with_args(obj.clone(), args) {
                        Ok(v) => v,
                        Err(e) => panic!("Myc Runtime Error (Closure): {}", e),
                    }
                } else {
                    res
                }
            }),
        })
    }

    fn specialize_node(&self, node: AnalyzedNode) -> AnalyzedNode {
        let registry = Rc::new(EnvFunctionRegistry {
            registry: self.function_registry.clone(),
            analyzed_registry: self.typed_function_registry.clone(),
        });

        let rtl_lookup = Rc::new(|name: &str, args: &[StaticType]| intrinsics::lookup(name, args));
        let typed_reg = self.typed_function_registry.clone();
        let syntax_reg = self.function_registry.clone();
        let mono_cache = self.monomorph_cache.clone();
        let root_values = self.root_values.clone();
        let root_types = self.root_types.clone();
        let root_purity = self.root_purity.clone();
        let optimization = self.optimization;

        let compiler = Rc::new(
            move |func_template: Rc<Node>,
                  arg_types: &[StaticType]|
                  -> Result<(Value, StaticType), String> {
                let mut diag = Diagnostics::new();
                let checker = TypeChecker::new(root_types.clone());
                let retyped_ast = checker.check(func_template.as_ref(), arg_types, &mut diag);

                if diag.has_errors() {
                    return Err(diag
                        .items
                        .into_iter()
                        .map(|d| d.message)
                        .collect::<Vec<_>>()
                        .join("\n"));
                }

                let analyzed = Analyzer::analyze(&retyped_ast, &root_purity.borrow());

                let sub_registry = Rc::new(EnvFunctionRegistry {
                    registry: syntax_reg.clone(),
                    analyzed_registry: typed_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(analyzed);

                let optimizer = Optimizer::new(optimization)
                    .with_globals(root_values.clone())
                    .with_purity(root_purity.clone());
                let optimized_ast = optimizer.optimize(specialized_ast);

                let exec_ast = Lowering::lower(optimized_ast);
                let mut vm = VM::new(root_values.clone());
                let compiled_val = match vm.run(&exec_ast) {
                    Ok(v) => v,
                    Err(e) => return Err(format!("VM Error during specialization: {}", e)),
                };

                let ret_type = exec_ast.ty.ty.clone();
                Ok((compiled_val, ret_type))
            },
        );

        let specializer = Specializer::new(
            Some(registry),
            Some(compiler),
            Some(rtl_lookup),
            Some(self.monomorph_cache.clone()),
        );

        specializer.specialize(node)
    }

    pub fn run_script(&self, source: &str) -> Result<Value, String> {
        self.preload_dependencies(source, None)?;
        if self.debug_mode {
            let (res, logs) = self.run_debug(source)?;
            for line in logs {
                println!("{}", line);
            }
            res
        } else {
            self.compile(source)
                .into_result()
                .and_then(|ast| self.run_script_compiled(ast))
        }
    }

    pub fn run_script_compiled(&self, compiled: TypedNode) -> Result<Value, String> {
        let linked = self.link(compiled);
        let func = self.instantiate(linked);
        let res = (func.func)(&[]);
        self.run_pipeline();
        Ok(res)
    }

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

        let mut vm = VM::new(self.root_values.clone());
        let mut observer = TracingObserver::new();

        // 1. Run the script wrapper (returns a closure representing the script)
        let result = vm.run_with_observer(&mut observer, &linked);
        // 2. Execute the root closure immediately to get the actual script result.
        // All Myc scripts are wrapped in a parameterless lambda for consistency.
        let mut final_result = result;
        if let Ok(Value::Object(obj)) = &final_result
            && let Some(_closure) = obj.as_any().downcast_ref::<crate::ast::vm::Closure>()
        {
            final_result = vm.run_with_args_observed(&mut observer, obj.clone(), &[]);
        }

        self.run_pipeline();

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