RustAst/src/ast/environment.rs

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

use crate::ast::nodes::{
    Address, AnalyzedNode, ExecNode, GlobalAnalyzedRegistry, GlobalFunctionRegistry, GlobalIdx,
    LambdaBinding, Node, NodeKind, VirtualId,
};
use crate::ast::compiler::dumper::Dumper;
use crate::ast::compiler::lambda_collector::LambdaCollector;
use crate::ast::compiler::module_loader::ModuleLoader;
use crate::ast::compiler::lowering::Lowering;
use crate::ast::compiler::macros::{MacroEvaluator, MacroExpander, MacroRegistry};
use crate::ast::compiler::optimizer::Optimizer;
use crate::ast::compiler::specializer::{FunctionRegistry, MonoCache, RtlLookupFunc, Specializer};
use crate::ast::rtl::{self, intrinsics};
use crate::ast::rtl::docs::{PipelineGenerator, RtlDocEntry, RtlRegistration};
use crate::ast::types::{
    CommentLine, NativeFunction, NodeIdentity, Purity, SourceLocation, StaticType, Value,
};

use crate::ast::diagnostics::Diagnostics;
pub use crate::ast::compiler::CompilationResult;

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

fn make_rtl_lookup() -> RtlLookupFunc {
    Rc::new(|name: &str, args: &[StaticType]| intrinsics::lookup(name, args))
}

/// Frozen snapshot of the RTL bootstrap state.
///
/// Created once after `rtl::register()` completes. All fields are immutable.
/// Multiple [`Environment`] instances can be created from a single `Rtl`
/// via [`Environment::from_rtl`], enabling independent execution contexts
/// that share the same native function set.
pub struct Rtl {
    /// The frozen RTL scope (scope index 0) containing all native bindings.
    pub scope: CompilerScope,
    /// Number of RTL slots — marks the boundary between RTL and user slots
    /// in the global value vectors.
    pub slot_count: u32,
    pub types: Vec<StaticType>,
    pub purity: Vec<Purity>,
    pub values: Rc<[Value]>,
    pub rtl_docs: Vec<RtlDocEntry>,
}

pub struct Environment {
    pub root_types: Rc<RefCell<Vec<StaticType>>>,
    pub root_purity: Rc<RefCell<Vec<Purity>>>,
    /// Frozen RTL value slice — shared across all environments from the same [`Rtl`].
    /// Empty placeholder during the bootstrap phase; set after the RTL freeze.
    rtl_values: Rc<[Value]>,
    /// Mutable user-defined global slots. Indexed as `[0..)` i.e. offset from `rtl.slot_count`.
    user_values: Rc<RefCell<Vec<Value>>>,
    pub fixed_scope_idx: i32,
    pub root_scopes: Rc<RefCell<Vec<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>>>,
    pub rtl_docs: Rc<RefCell<Vec<RtlDocEntry>>>,
    /// Documentation for symbols defined in Myc source (via `;;` comments before `def`/`macro`).
    /// Key: symbol name. Value: concatenated doc lines joined by newlines.
    pub myc_docs: Rc<RefCell<HashMap<String, String>>>,
    /// Frozen RTL snapshot from which this environment was created.
    pub rtl: Rc<Rtl>,
}

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

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

struct RuntimeMacroEvaluator {
    root_scopes: Rc<RefCell<Vec<CompilerScope>>>,
    root_slot_count: Rc<RefCell<u32>>,
    root_types: Rc<RefCell<Vec<StaticType>>>,
    globals: GlobalStore,
    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 { symbol: sym, .. } = &node.kind
            && let Some(arg_node) = bindings.get(&sym.name)
        {
            return Ok(Value::Quote(Rc::new(arg_node.clone())));
        }

        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 = 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.format_errors());
        }

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

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

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

impl Environment {
    pub fn new() -> Self {
        // Phase 1: Bootstrap — fixed_scope_idx = -1 allows allocate_slot to
        // write freely into scope 0 (the RTL scope).
        let mut env = Self {
            root_types: Rc::new(RefCell::new(Vec::new())),
            root_purity: Rc::new(RefCell::new(Vec::new())),
            rtl_values: Rc::from([]),   // placeholder — filled after freeze
            user_values: Rc::new(RefCell::new(Vec::new())), // bootstrap scratch during RTL phase
            fixed_scope_idx: -1,
            root_scopes: Rc::new(RefCell::new(vec![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())),
            rtl_docs: Rc::new(RefCell::new(Vec::new())),
            myc_docs: Rc::new(RefCell::new(HashMap::new())),
            // Placeholder — overwritten below after bootstrap.
            rtl: Rc::new(Rtl {
                scope: CompilerScope::new(),
                slot_count: 0,
                types: vec![],
                purity: vec![],
                values: Rc::from([]),
                rtl_docs: vec![],
            }),
        };

        rtl::register(&env);

        // Phase 2: Freeze scope 0 and snapshot the RTL state.
        env.fixed_scope_idx = 0;
        let rtl_slot_count = *env.root_slot_count.borrow();
        // user_values was used as the bootstrap scratch; freeze it into an immutable slice.
        let rtl_vals: Rc<[Value]> = env.user_values.borrow().clone().into();
        env.rtl = Rc::new(Rtl {
            scope: env.root_scopes.borrow()[0].clone(),
            slot_count: rtl_slot_count,
            types: env.root_types.borrow().clone(),
            purity: env.root_purity.borrow().clone(),
            values: Rc::clone(&rtl_vals),
            rtl_docs: std::mem::take(&mut env.rtl_docs.borrow_mut()),
        });

        // Set the frozen RTL slice. Stream closures registered during bootstrap
        // (e.g. in streams.rs) hold a GlobalStore that captured the OLD user_values
        // Rc (bootstrap scratch = RTL values, never written again). After resetting
        // user_values here, script VMs use the new empty Rc and grow it with user
        // slots. Streams cannot access user globals — RTL-only invariant enforced.
        env.rtl_values = Rc::clone(&rtl_vals);
        env.user_values = Rc::new(RefCell::new(Vec::new()));

        // Push the first mutable user scope (Level 1)
        env.root_scopes.borrow_mut().push(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
    }

    /// Creates a fresh execution environment from a frozen [`Rtl`] snapshot.
    ///
    /// The new environment starts with the RTL slots pre-loaded and a clean
    /// user scope. Multiple environments created from the same `Rtl` are fully
    /// independent: user definitions, macros, and loaded modules do not leak
    /// across them.
    pub fn from_rtl(rtl: Rc<Rtl>) -> Self {
        let scopes = vec![rtl.scope.clone(), CompilerScope::new()];

        let env = Self {
            root_types: Rc::new(RefCell::new(rtl.types.clone())),
            root_purity: Rc::new(RefCell::new(rtl.purity.clone())),
            rtl_values: Rc::clone(&rtl.values), // O(1) Rc increment — no clone of RTL values!
            user_values: Rc::new(RefCell::new(Vec::new())),
            fixed_scope_idx: 0,
            root_scopes: Rc::new(RefCell::new(scopes)),
            root_slot_count: Rc::new(RefCell::new(rtl.slot_count)),
            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())),
            rtl_docs: Rc::new(RefCell::new(Vec::new())),
            myc_docs: Rc::new(RefCell::new(HashMap::new())),
            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
    }

    /// Returns the frozen RTL snapshot this environment was bootstrapped from.
    pub fn rtl(&self) -> Rc<Rtl> {
        Rc::clone(&self.rtl)
    }

    /// Returns a [`GlobalStore`] view over this environment's global value space.
    ///
    /// The RTL portion is a shared immutable slice; the user portion is a
    /// per-environment mutable vec. Both are reference-counted — cloning is cheap.
    pub fn global_store(&self) -> GlobalStore {
        GlobalStore::new(
            Rc::clone(&self.rtl_values),
            Rc::clone(&self.user_values),
            self.rtl.slot_count as usize,
        )
    }

    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(),
            globals: self.global_store(),
            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.
    /// Used to pre-load all dependencies of a script before compiling it.
    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 user dependencies (file I/O + parse, topological order)
        let loader = ModuleLoader::new(
            Rc::clone(&self.search_paths),
            Rc::clone(&self.loaded_modules),
        );
        files.extend(loader.collect_dependency_files(source, base_path)?);

        // 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 { pattern, .. } => {
                let mut root_scopes = self.root_scopes.borrow_mut();
                let last_idx = root_scopes.len() - 1;
                let current_scope = &mut root_scopes[last_idx];

                fn register_pattern(
                    pattern: &SyntaxNode,
                    scope: &mut CompilerScope,
                    slot_count: &mut u32,
                    identity: &crate::ast::types::NodeIdentity,
                ) {
                    match &pattern.kind {
                        SyntaxKind::Identifier { symbol: sym, .. } => {
                            if !scope.locals.contains_key(sym) {
                                let slot = VirtualId(*slot_count);
                                scope.locals.insert(
                                    sym.clone(),
                                    LocalInfo {
                                        addr: Address::Local(slot),
                                        identity: Rc::new(identity.clone()),
                                        _ty: StaticType::Any,
                                        purity: Purity::Impure,
                                    },
                                );
                                *slot_count += 1;
                            }
                        }
                        SyntaxKind::Tuple { elements } => {
                            for el in elements {
                                register_pattern(el, scope, slot_count, identity);
                            }
                        }
                        _ => {}
                    }
                }

                let mut slot_count = self.root_slot_count.borrow_mut();
                register_pattern(pattern, current_scope, &mut slot_count, &node.identity);
            }
            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 { symbol: sym, .. } => vec![sym.name.clone()],
                        SyntaxKind::Tuple { elements } => {
                            elements.iter().flat_map(|e| extract_names(e)).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);
                }
            }
            SyntaxKind::Expansion { expanded, .. } => {
                self.discover_globals(expanded);
            }
            _ => {}
        }
    }

    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 = CapturePass::apply(bound_ast, &captures);

        // Pre-allocate user global slots to prevent out-of-bounds during specialization/optimization.
        // root_slot_count includes RTL slots; subtract to get the number of user-only slots.
        {
            let mut user = self.user_values.borrow_mut();
            let count = *self.root_slot_count.borrow() as usize;
            let user_count = count.saturating_sub(self.rtl.slot_count as usize);
            if user_count > user.len() {
                user.resize(user_count, 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 NodeKind::Lambda { .. } = bound_ast.kind {
            bound_ast
        } else {
            Node {
                identity: bound_ast.identity.clone(),
                kind: NodeKind::Lambda {
                    params: std::rc::Rc::new(Node {
                        identity: bound_ast.identity.clone(),
                        kind: NodeKind::Tuple { elements: vec![] },
                        ty: (),
                        comments: Rc::from([]),
                    }),
                    body: std::rc::Rc::new(bound_ast),
                    info: LambdaBinding {
                        upvalues: vec![],
                        positional_count: Some(0),
                    },
                },
                ty: (),
                comments: Rc::from([]),
            }
        };
        let typed = checker.check(&wrapped_ast, &[], diagnostics);
        self.collect_doc_comments(&typed);
        Some(typed)
    }

    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.format_errors());
        }

        Ok(typed_ast_opt.unwrap())
    }

    /// Allocates a new global slot and registers a value in the fixed RTL scope (scope index 0).
    /// This is the single source of truth for all global slot allocation.
    fn allocate_slot(&self, name: &str, ty: StaticType, purity: Purity, val: Value) {
        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),
            LocalInfo {
                addr: Address::Global(global_idx),
                identity: NodeIdentity::new(SourceLocation { line: 0, col: 0 }),
                _ty: ty.clone(),
                purity,
            },
        );
        *slot_count += 1;
        self.root_types.borrow_mut().push(ty);
        self.root_purity.borrow_mut().push(purity);
        self.user_values.borrow_mut().push(val);
    }

    pub fn register_native(
        &self,
        name: &str,
        ty: StaticType,
        func: Rc<NativeFunction>,
    ) {
        self.allocate_slot(name, ty, func.purity, Value::Function(func));
    }

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

    pub fn register_constant(&self, name: &str, ty: StaticType, val: Value) -> RtlRegistration {
        self.allocate_slot(name, ty, Purity::Pure, val);
        RtlRegistration::new(name.to_string(), Rc::clone(&self.rtl_docs))
    }


    /// Returns the names of all bindings registered in the fixed RTL scope.
    /// Useful for introspection (e.g., MCP server listing available built-ins).
    pub fn list_bindings(&self) -> Vec<String> {
        let mut names: Vec<String> = self.rtl.scope
            .locals
            .keys()
            .map(|sym| sym.name.to_string())
            .collect();
        names.sort();
        names
    }

    /// Returns formatted documentation for all RTL symbols that have been annotated
    /// via the `.doc()` builder. Each entry includes the name, type signature
    /// (derived from `StaticType`), one-liner, optional description, and examples.
    pub fn list_rtl_docs(&self) -> Vec<String> {
        use crate::ast::nodes::Symbol;
        let scope = &self.rtl.scope;

        let mut entries: Vec<String> = self.rtl.rtl_docs.iter().map(|entry| {
            // Derive the signature from the registered StaticType.
            let sig = scope
                .locals
                .get(&Symbol::from(entry.name.as_str()))
                .map(|info| info._ty.to_doc_string())
                .unwrap_or_else(|| "unknown".to_string());

            let mut out = format!("{} : {}\n  {}", entry.name, sig, entry.one_liner);

            if let Some(desc) = entry.description {
                out.push_str(&format!("\n  {}", desc));
            }
            if let Some(examples) = entry.examples {
                out.push_str("\n  Examples:");
                for ex in examples.iter() {
                    out.push_str(&format!("\n    {}", ex));
                }
            }
            out
        }).collect();

        // Include Myc-defined symbol docs (from `;;` comments in source)
        let myc_docs = self.myc_docs.borrow();
        for (name, doc_text) in myc_docs.iter() {
            let sig = scope
                .locals
                .get(&Symbol::from(name.as_str()))
                .map(|info| info._ty.to_doc_string())
                .unwrap_or_else(|| "unknown".to_string());
            entries.push(format!("{} : {}\n  {}", name, sig, doc_text));
        }

        entries.sort_by_key(|a| a.to_lowercase());
        entries
    }

    /// Returns formatted documentation for a single RTL symbol by name,
    /// or `None` if the symbol has no doc entry.
    pub fn get_rtl_doc(&self, name: &str) -> Option<String> {
        use crate::ast::nodes::Symbol;
        let sig = self.rtl.scope
            .locals
            .get(&Symbol::from(name))
            .map(|info| info._ty.to_doc_string())
            .unwrap_or_else(|| "unknown".to_string());

        // Check RTL docs first
        if let Some(entry) = self.rtl.rtl_docs.iter().find(|e| e.name == name) {
            let mut out = format!("{} : {}\n  {}", entry.name, sig, entry.one_liner);
            if let Some(desc) = entry.description {
                out.push_str(&format!("\n  {}", desc));
            }
            if let Some(examples) = entry.examples {
                out.push_str("\n  Examples:");
                for ex in examples.iter() {
                    out.push_str(&format!("\n    {}", ex));
                }
            }
            return Some(out);
        }

        // Fall back to Myc-source docs
        let myc_docs = self.myc_docs.borrow();
        myc_docs
            .get(name)
            .map(|doc_text| format!("{} : {}\n  {}", name, sig, doc_text))
    }

    /// Walks a typed AST and registers `;;` doc comments on `def`/`macro` nodes
    /// whose target is a plain identifier into the `myc_docs` registry.
    fn collect_doc_comments(&self, node: &TypedNode) {
        // Extract the symbol name this node defines, if any.
        let sym_name: Option<Rc<str>> = match &node.kind {
            NodeKind::Def { pattern, .. } => {
                if let NodeKind::Identifier { symbol, .. } = &pattern.kind {
                    Some(symbol.name.clone())
                } else {
                    None
                }
            }
            NodeKind::MacroDecl { name, .. } => Some(name.name.clone()),
            NodeKind::Block { exprs } => {
                for expr in exprs {
                    self.collect_doc_comments(expr);
                }
                return;
            }
            NodeKind::Lambda { body, .. } => {
                self.collect_doc_comments(body);
                return;
            }
            _ => None,
        };

        if let Some(name) = sym_name {
            let doc_text: String = node
                .comments
                .iter()
                .filter_map(|line| match line {
                    CommentLine::Doc(text) => Some(text.as_ref()),
                    _ => None,
                })
                .collect::<Vec<_>>()
                .join("\n");

            if !doc_text.is_empty() {
                self.myc_docs.borrow_mut().insert(name.to_string(), doc_text);
            }
        }
    }

    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.global_store())
            .with_purity(self.root_purity.clone())
            .with_registry(self.typed_function_registry.clone());
        let optimized = optimizer.optimize(specialized);

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

    /// Converts a linked `ExecNode` into a callable `Closure`.
    ///
    /// Fast path: if the node is already a top-level lambda without upvalues,
    /// the `Closure` is constructed directly without running the VM.
    ///
    /// Slow path: the node is executed once to obtain its resulting `Closure`
    /// value (e.g. a `do`-block that returns a lambda).
    ///
    /// The caller is responsible for creating a `VM` and invoking the closure
    /// via `vm.run_with_args`. This keeps VM lifecycle and error handling at
    /// the call site, where the required strategy (single run vs. repeated
    /// benchmark iterations) is known.
    pub fn instantiate(&self, node: ExecNode) -> Result<Rc<Closure>, String> {
        // Fast path: top-level lambda without upvalues — build Closure directly.
        if let NodeKind::Lambda { params, body, info } = &node.kind
            && info.upvalues.is_empty()
        {
            return Ok(Rc::new(Closure::new(
                params.clone(),
                body.ty.original.clone(),
                body.clone(),
                Vec::new(),
                info.positional_count,
                node.ty.stack_size,
            )));
        }

        // Slow path: run the node once to extract the resulting Closure.
        let mut vm = VM::new(self.global_store());
        let res = vm.run(&node)?;
        if let Value::Closure(obj) = res {
            Ok(obj)
        } else {
            Err("Script did not produce a callable closure".to_string())
        }
    }

    /// Creates a new `VM` connected to this environment's global scope.
    pub fn create_vm(&self) -> VM {
        VM::new(self.global_store())
    }

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

        let rtl_lookup = make_rtl_lookup();
        let typed_reg = self.typed_function_registry.clone();
        let mono_cache = self.monomorph_cache.clone();
        let root_values = self.global_store();
        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<AnalyzedPhase>>,
                  arg_types: &[StaticType]|
                  -> Result<(Value, StaticType), String> {
                let mut diag = Diagnostics::new();
                let checker = TypeChecker::new(root_types.clone());

                // Monomorphization: re-type-check the function template with the concrete
                // call-site argument types, producing a specialized TypedNode.
                //
                // `func_template` is an AnalyzedNode whose `ty.original` holds the Rc<TypedNode>
                // preserved by the Analyzer in NodeMetrics.
                //
                // `check_node_as_bound` is generic over all `BoundLike` phases (BoundPhase,
                // TypedPhase, and AnalyzedPhase all share the same binding structure). This lets
                // us re-run type inference on the TypedNode with concrete `arg_types`: existing
                // `ty` metadata is discarded and all types are inferred fresh from the call site.
                let retyped_ast = checker.check_node_as_bound(func_template.ty.original.as_ref(), arg_types, &mut diag);

                if diag.has_errors() {
                    return Err(diag.format_errors());
                }

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

                let sub_registry = Rc::new(EnvFunctionRegistry {
                    analyzed_registry: typed_reg.clone(),
                });
                let sub_rtl_lookup = make_rtl_lookup();

                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 closure = self.instantiate(linked)?;
        let mut vm = self.create_vm();
        let res = vm.run_with_args(closure, &[])?;
        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.global_store());
        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::Closure(obj)) = &final_result {
            final_result = vm.run_with_args_observed(&mut observer, obj.clone(), &[]);
        }

        self.run_pipeline();

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