//! AST → bytecode compiler for the Aster VM.
//!
//! Walks the AST recursively, emitting stack-based bytecode instructions.
//! Handles local variable resolution (slot indices), upvalue capture for
//! closures, jump backpatching for control flow, and constant pool management.

use crate::ast::*;
use crate::ast::expr::{Literal, UnaryOp, BinaryOp, LogicalOp, AssignOp};
use crate::error::RuntimeError;
use crate::runtime::{Value, FunctionProto};
use super::opcode::*;

use std::rc::Rc;
use std::cell::RefCell;

// ============================================================================
// Compiler
// ============================================================================

#[derive(Clone)]
struct Local {
    name: String,
    depth: u8,       // scope depth where declared; 0 = uninitialized
    is_captured: bool,
    is_const: bool,
}

#[derive(Clone)]
struct Upvalue {
    index: u8,
    is_local: bool,  // true = captured from enclosing fn's local; false = from upvalue
    is_const: bool,  // true = the source variable was declared `const`
    name: String,    // variable name (for transitive upvalue resolution)
}

struct LoopContext {
    break_patches: Vec<usize>,
    continue_patches: Vec<usize>,
}

pub struct Compiler {
    function: FunctionProto,
    locals: Vec<Local>,
    /// Upvalues for this function (shared with child for transitive resolution)
    upvalues: Rc<RefCell<Vec<Upvalue>>>,
    /// ALL enclosing locals from entire chain (merged, direct parent first)
    enclosing_locals: Option<Vec<Local>>,
    /// Direct parent's own locals count (to distinguish parent locals from deeper ones)
    parent_locals_count: usize,
    /// Direct parent's upvalues (shared)
    enclosing_upvalues: Option<Rc<RefCell<Vec<Upvalue>>>>,
    /// Grandparent's upvalues (shared, for 3-level transitive capture)
    grandparent_upvalues: Option<Rc<RefCell<Vec<Upvalue>>>>,
    scope_depth: u8,
    loop_stack: Vec<LoopContext>,
}

impl Compiler {
    pub fn new(name: Option<String>) -> Self {
        Self {
            function: FunctionProto::new(name),
            locals: Vec::new(),
            upvalues: Rc::new(RefCell::new(Vec::new())),
            enclosing_locals: None,
            parent_locals_count: 0,
            enclosing_upvalues: None,
            grandparent_upvalues: None,
            scope_depth: 0,
            loop_stack: Vec::new(),
        }
    }

    /// Convenience: compile a list of statements into a FunctionProto.
    pub fn compile(stmts: &[Stmt]) -> Result<FunctionProto, RuntimeError> {
        let mut compiler = Self::new(None);
        for stmt in stmts {
            compiler.compile_stmt(stmt)?;
        }
        // Implicit return nil at end of function
        compiler.emit_op(OpCode::LoadNil);
        compiler.emit_op(OpCode::Return);
        Ok(compiler.function)
    }

    /// Compile a list of statements (for use from nested compilers).
    fn compile_stmts(&mut self, stmts: &[Stmt]) -> Result<(), RuntimeError> {
        for stmt in stmts {
            self.compile_stmt(stmt)?;
        }
        Ok(())
    }

    // ========================================================================
    // Statement compilation
    // ========================================================================

    fn compile_stmt(&mut self, stmt: &Stmt) -> Result<(), RuntimeError> {
        match stmt {
            Stmt::Let { name, initializer, mutable } => {
                self.compile_let(name, initializer, *mutable)?;
            }
            Stmt::ExprStmt(expr) => {
                self.compile_expr(expr)?;
                self.emit_op(OpCode::Pop);
            }
            Stmt::Block(stmts) => {
                self.begin_scope();
                self.compile_stmts(stmts)?;
                self.end_scope();
            }
            Stmt::If { condition, then_branch, else_branch } => {
                self.compile_if(condition, then_branch, else_branch.as_deref())?;
            }
            Stmt::While { condition, body } => {
                self.compile_while(condition, body)?;
            }
            Stmt::For { initializer, condition, step, body } => {
                self.compile_for(initializer.as_deref(), condition.as_ref(), step.as_ref(), body)?;
            }
            Stmt::ForIn { var_name, iterable, body } => {
                self.compile_for_in(var_name, iterable, body)?;
            }
            Stmt::Function { name, params, body } => {
                self.compile_function_decl(name, params, body)?;
                // DefineGlobal pushes the value back; pop it as this is a statement
                self.emit_op(OpCode::Pop);
            }
            Stmt::Return(expr) => {
                if let Some(e) = expr {
                    self.compile_expr(e)?;
                } else {
                    self.emit_op(OpCode::LoadNil);
                }
                self.emit_op(OpCode::Return);
            }
            Stmt::Break => {
                self.compile_break()?;
            }
            Stmt::Continue => {
                self.compile_continue()?;
            }
        }
        Ok(())
    }

    fn compile_let(&mut self, name: &str, initializer: &Expr, mutable: bool) -> Result<(), RuntimeError> {
        self.compile_expr(initializer)?;
        if self.scope_depth == 0 {
            let name_idx = self.add_string_constant(name);
            emit_u16(&mut self.function.code, OpCode::DefineGlobal, name_idx);
            self.function.code.push(if mutable { 1 } else { 0 });
            // DefineGlobal pushes value back; pop it for statement-level let
            self.emit_op(OpCode::Pop);
        } else {
            let slot = self.locals.len() as u8;
            self.locals.push(Local {
                name: name.to_string(),
                depth: self.scope_depth,
                is_captured: false,
                is_const: !mutable,
            });
            // StoreLocal PEEKS the value — it stays on stack as the local
            emit_u8(&mut self.function.code, OpCode::StoreLocal, slot);
        }
        Ok(())
    }

    fn compile_if(&mut self, condition: &Expr, then_branch: &Stmt, else_branch: Option<&Stmt>) -> Result<(), RuntimeError> {
        self.compile_expr(condition)?;
        let else_jump = self.emit_jump(OpCode::JumpIfFalse);
        self.compile_stmt(then_branch)?;
        if let Some(else_stmt) = else_branch {
            let end_jump = self.emit_jump(OpCode::Jump);
            self.patch_jump(else_jump);
            self.compile_stmt(else_stmt)?;
            self.patch_jump(end_jump);
        } else {
            self.patch_jump(else_jump);
        }
        Ok(())
    }

    fn compile_while(&mut self, condition: &Expr, body: &Stmt) -> Result<(), RuntimeError> {
        let start_ip = self.function.code.len();
        self.compile_expr(condition)?;
        let exit_jump = self.emit_jump(OpCode::JumpIfFalse);

        self.loop_stack.push(LoopContext {
            break_patches: Vec::new(),
            continue_patches: Vec::new(),
        });

        self.compile_stmt(body)?;
        self.emit_loop_jump(start_ip);

        self.patch_jump(exit_jump);

        let loop_ctx = self.loop_stack.pop().unwrap();
        for patch in loop_ctx.break_patches {
            self.patch_jump(patch);
        }
        for patch in loop_ctx.continue_patches {
            self.patch_jump_to(patch, start_ip);
        }
        Ok(())
    }

    fn compile_for(&mut self, initializer: Option<&Stmt>, condition: Option<&Expr>, step: Option<&Expr>, body: &Stmt) -> Result<(), RuntimeError> {
        // For loop: new scope for init
        self.begin_scope();

        if let Some(init) = initializer {
            self.compile_stmt(init)?;
        }

        let start_ip = self.function.code.len();
        if let Some(cond) = condition {
            self.compile_expr(cond)?;
        } else {
            self.emit_op(OpCode::LoadTrue);
        }
        let exit_jump = self.emit_jump(OpCode::JumpIfFalse);

        self.loop_stack.push(LoopContext {
            break_patches: Vec::new(),
            continue_patches: Vec::new(),
        });

        self.compile_stmt(body)?;

        // continue lands here (after body, before step)
        let continue_ip = self.function.code.len();
        let loop_ctx = self.loop_stack.pop().unwrap();
        for patch in loop_ctx.continue_patches {
            self.patch_jump_to(patch, continue_ip);
        }

        if let Some(s) = step {
            self.compile_expr(s)?;
            self.emit_op(OpCode::Pop);
        }

        self.emit_loop_jump(start_ip);
        self.patch_jump(exit_jump);

        for patch in loop_ctx.break_patches {
            self.patch_jump(patch);
        }

        self.end_scope();
        Ok(())
    }

    fn compile_for_in(&mut self, var_name: &str, iterable: &Expr, body: &Stmt) -> Result<(), RuntimeError> {
        // 1. New scope for hidden iterator locals (items, idx)
        self.begin_scope();

        // 2. Evaluate iterable and compute items array
        self.compile_expr(iterable)?;
        emit_op(&mut self.function.code, OpCode::ForInInit); // pops iterable, pushes items[]

        // 3. Store items array in a hidden local (peek, no Pop — cleaned by end_scope)
        let items_slot = self.locals.len() as u8;
        self.locals.push(Local {
            name: format!("__iter_items_{}", items_slot),
            depth: self.scope_depth,
            is_captured: false,
            is_const: false,
        });
        emit_u8(&mut self.function.code, OpCode::StoreLocal, items_slot);

        // 4. Initialize index = 0 in hidden local
        let idx_slot = self.locals.len() as u8;
        self.locals.push(Local {
            name: format!("__iter_idx_{}", idx_slot),
            depth: self.scope_depth,
            is_captured: false,
            is_const: false,
        });
        let zero_idx = self.add_constant(Value::Number(0.0));
        emit_u16(&mut self.function.code, OpCode::LoadConst, zero_idx);
        emit_u8(&mut self.function.code, OpCode::StoreLocal, idx_slot);

        // 5. ForInNext: reads items[idx_slot], idx[idx_slot], pushes element
        let forin_loc = self.emit_forin_next(items_slot, idx_slot);

        // 6. Store loop variable in a local (peek, value stays on stack)
        let loop_var_slot = self.locals.len() as u8;
        self.locals.push(Local {
            name: var_name.to_string(),
            depth: self.scope_depth,
            is_captured: false,
            is_const: false,
        });
        emit_u8(&mut self.function.code, OpCode::StoreLocal, loop_var_slot);

        // 7. Loop body (can resolve var_name to loop_var_slot)
        self.loop_stack.push(LoopContext {
            break_patches: Vec::new(),
            continue_patches: Vec::new(),
        });

        self.compile_stmt(body)?;

        // Patch continue
        let continue_ip = self.function.code.len();
        let loop_ctx = self.loop_stack.pop().unwrap();
        for patch in loop_ctx.continue_patches {
            self.patch_jump_to(patch, continue_ip);
        }

        // 8. Pop loop var element and remove from tracking
        self.emit_op(OpCode::Pop);
        self.locals.pop(); // loop_var_slot

        // 9. Jump back to ForInNext (next iteration re-adds loop var via StoreLocal)
        self.emit_loop_jump(forin_loc);

        // 10. Exit target: patch ForInNext exit and break jumps
        let exit_ip = self.function.code.len();
        self.patch_forin_jump(forin_loc, exit_ip);
        for patch in loop_ctx.break_patches {
            self.patch_jump_to(patch, exit_ip);
        }

        // 11. Pop loop var (ForInNext pushes Nil on exit; break leaves element on stack)
        self.emit_op(OpCode::Pop);

        // 12. End scope: pops items_slot + idx_slot values (left by StoreLocal peek)
        self.end_scope();
        Ok(())
    }

    fn compile_function_decl(&mut self, name: &str, params: &[String], body: &[Stmt]) -> Result<(), RuntimeError> {
        let proto = self.compile_nested_function(Some(name.to_string()), params, body)?;
        let upvalues = proto.upvalues.clone();
        let proto_idx = self.add_function_proto_constant(proto);
        self.emit_closure(proto_idx, &upvalues);

        // Bind to name
        if self.scope_depth == 0 {
            let name_idx = self.add_string_constant(name);
            emit_u16(&mut self.function.code, OpCode::DefineGlobal, name_idx);
            self.function.code.push(1); // mutable=true for fn declarations
        } else {
            let slot = self.locals.len() as u8;
            self.locals.push(Local {
                name: name.to_string(),
                depth: self.scope_depth,
                is_captured: false,
                is_const: false,
            });
            emit_u8(&mut self.function.code, OpCode::StoreLocal, slot);
        }
        Ok(())
    }

    fn compile_break(&mut self) -> Result<(), RuntimeError> {
        let jump_loc = self.emit_jump_placeholder();
        let loop_ctx = self.loop_stack.last_mut().ok_or_else(|| RuntimeError::RuntimeError {
            message: "break outside of loop".into(),
            token: None,
        })?;
        loop_ctx.break_patches.push(jump_loc);
        Ok(())
    }

    fn compile_continue(&mut self) -> Result<(), RuntimeError> {
        let jump_loc = self.emit_jump_placeholder();
        let loop_ctx = self.loop_stack.last_mut().ok_or_else(|| RuntimeError::RuntimeError {
            message: "continue outside of loop".into(),
            token: None,
        })?;
        loop_ctx.continue_patches.push(jump_loc);
        Ok(())
    }

    // ========================================================================
    // Expression compilation
    // ========================================================================

    fn compile_expr(&mut self, expr: &Expr) -> Result<(), RuntimeError> {
        match expr {
            Expr::Literal(lit) => self.compile_literal(lit),
            Expr::Variable(name) => self.compile_variable(name),
            Expr::Assign { name, op, value } => self.compile_assign(name, op, value),
            Expr::Get { object, name } => self.compile_get(object, name),
            Expr::Set { object, name, op, value } => self.compile_set(object, name, op, value),
            Expr::ObjectLiteral { properties } => self.compile_object(properties),
            Expr::ArrayLiteral { elements } => self.compile_array(elements),
            Expr::IndexGet { array, index } => self.compile_index_get(array, index),
            Expr::IndexSet { array, index, op, value } => self.compile_index_set(array, index, op, value),
            Expr::Unary { op, right } => self.compile_unary(op, right),
            Expr::Binary { left, op, right } => self.compile_binary(left, op, right),
            Expr::Logical { left, op, right } => self.compile_logical(left, op, right),
            Expr::Ternary { condition, then_branch, else_branch } => {
                self.compile_ternary(condition, then_branch, else_branch)
            }
            Expr::Call { callee, arguments } => self.compile_call(callee, arguments),
            Expr::Lambda { params, body } => self.compile_lambda(params, body),
        }
    }

    fn compile_literal(&mut self, lit: &Literal) -> Result<(), RuntimeError> {
        match lit {
            Literal::Number(n) => {
                let idx = self.function.add_constant(Value::Number(*n));
                emit_u16(&mut self.function.code, OpCode::LoadConst, idx);
            }
            Literal::String(s) => {
                let idx = self.function.add_constant(Value::String(s.clone()));
                emit_u16(&mut self.function.code, OpCode::LoadConst, idx);
            }
            Literal::Bool(true) => self.emit_op(OpCode::LoadTrue),
            Literal::Bool(false) => self.emit_op(OpCode::LoadFalse),
            Literal::Nil => self.emit_op(OpCode::LoadNil),
        }
        Ok(())
    }

    fn compile_variable(&mut self, name: &str) -> Result<(), RuntimeError> {
        // Try to resolve as local
        if let Some(slot) = self.resolve_local(name) {
            emit_u8(&mut self.function.code, OpCode::LoadLocal, slot);
            return Ok(());
        }
        // Try to resolve as upvalue
        if let Some(upvalue_idx) = self.resolve_upvalue(name) {
            emit_u8(&mut self.function.code, OpCode::LoadUpvalue, upvalue_idx);
            return Ok(());
        }
        // Fall back to global
        let name_idx = self.add_string_constant(name);
        emit_u16(&mut self.function.code, OpCode::LoadGlobal, name_idx);
        Ok(())
    }

    fn compile_assign(&mut self, name: &str, op: &AssignOp, value: &Expr) -> Result<(), RuntimeError> {
        if *op == AssignOp::Equal {
            // Simple assignment
            self.compile_expr(value)?;
            if let Some(slot) = self.resolve_local(name) {
                // Check const for local
                if self.locals[slot as usize].is_const {
                    return Err(RuntimeError::RuntimeError {
                        message: format!("Cannot reassign constant '{}'", name),
                        token: None,
                    });
                }
                emit_u8(&mut self.function.code, OpCode::StoreLocal, slot);
            } else if let Some(uv_idx) = self.resolve_upvalue(name) {
                // Check const for upvalue
                let upvalues = self.upvalues.borrow();
                if upvalues[uv_idx as usize].is_const {
                    return Err(RuntimeError::RuntimeError {
                        message: format!("Cannot reassign constant '{}'", name),
                        token: None,
                    });
                }
                emit_u8(&mut self.function.code, OpCode::StoreUpvalue, uv_idx);
            } else {
                let name_idx = self.add_string_constant(name);
                emit_u16(&mut self.function.code, OpCode::StoreGlobal, name_idx);
            }
        } else {
            // Compound assignment
            self.compile_expr(value)?;
            let compound_op = match op {
                AssignOp::PlusEqual => CompoundOp::PlusEqual,
                AssignOp::MinusEqual => CompoundOp::MinusEqual,
                AssignOp::StarEqual => CompoundOp::StarEqual,
                AssignOp::SlashEqual => CompoundOp::SlashEqual,
                AssignOp::PercentEqual => CompoundOp::PercentEqual,
                AssignOp::Equal => unreachable!(),
            };
            if let Some(slot) = self.resolve_local(name) {
                emit_u8(&mut self.function.code, OpCode::CompoundAssignLocal, slot);
                self.function.code.push(compound_op as u8);
            } else if let Some(uv_idx) = self.resolve_upvalue(name) {
                emit_u8(&mut self.function.code, OpCode::CompoundAssignUpvalue, uv_idx);
                self.function.code.push(compound_op as u8);
            } else {
                let name_idx = self.add_string_constant(name);
                emit_u16(&mut self.function.code, OpCode::CompoundAssignGlobal, name_idx);
                self.function.code.push(compound_op as u8);
            }
        }
        Ok(())
    }

    fn compile_get(&mut self, object: &Expr, name: &str) -> Result<(), RuntimeError> {
        self.compile_expr(object)?;
        let name_idx = self.add_string_constant(name);
        emit_u16(&mut self.function.code, OpCode::GetProperty, name_idx);
        Ok(())
    }

    fn compile_set(&mut self, object: &Expr, name: &str, op: &AssignOp, value: &Expr) -> Result<(), RuntimeError> {
        if *op == AssignOp::Equal {
            self.compile_expr(object)?;
            self.compile_expr(value)?;
            let name_idx = self.add_string_constant(name);
            emit_u16(&mut self.function.code, OpCode::SetProperty, name_idx);
        } else {
            // Compound property set: object.name op= value
            // CompoundAssignProp handler does get_property internally
            self.compile_expr(object)?;
            self.compile_expr(value)?; // rhs
            let compound_op = assign_op_to_compound(op);
            let name_idx = self.add_string_constant(name);
            emit_u16(&mut self.function.code, OpCode::CompoundAssignProp, name_idx);
            self.function.code.push(compound_op as u8);
        }
        Ok(())
    }

    fn compile_object(&mut self, properties: &[(String, Expr)]) -> Result<(), RuntimeError> {
        emit_op(&mut self.function.code, OpCode::NewObject);
        for (key, value_expr) in properties {
            self.emit_op(OpCode::Dup);
            self.compile_expr(value_expr)?;
            let name_idx = self.add_string_constant(key);
            emit_u16(&mut self.function.code, OpCode::SetProperty, name_idx);
            // SetProperty leaves the value on stack; we need the object, so Pop the value
            self.emit_op(OpCode::Pop);
        }
        Ok(())
    }

    fn compile_array(&mut self, elements: &[Expr]) -> Result<(), RuntimeError> {
        for e in elements {
            self.compile_expr(e)?;
        }
        emit_u16(&mut self.function.code, OpCode::NewArray, elements.len() as u16);
        Ok(())
    }

    fn compile_index_get(&mut self, array: &Expr, index: &Expr) -> Result<(), RuntimeError> {
        self.compile_expr(array)?;
        self.compile_expr(index)?;
        emit_op(&mut self.function.code, OpCode::GetIndex);
        Ok(())
    }

    fn compile_index_set(&mut self, array: &Expr, index: &Expr, op: &AssignOp, value: &Expr) -> Result<(), RuntimeError> {
        self.compile_expr(array)?;
        self.compile_expr(index)?;
        self.compile_expr(value)?;
        if *op == AssignOp::Equal {
            emit_op(&mut self.function.code, OpCode::SetIndex);
        } else {
            let compound_op = assign_op_to_compound(op);
            emit_u8(&mut self.function.code, OpCode::CompoundAssignIndex, compound_op as u8);
        }
        Ok(())
    }

    fn compile_unary(&mut self, op: &UnaryOp, right: &Expr) -> Result<(), RuntimeError> {
        self.compile_expr(right)?;
        match op {
            UnaryOp::Negate => self.emit_op(OpCode::Negate),
            UnaryOp::Not => self.emit_op(OpCode::Not),
        }
        Ok(())
    }

    fn compile_binary(&mut self, left: &Expr, op: &BinaryOp, right: &Expr) -> Result<(), RuntimeError> {
        self.compile_expr(left)?;
        self.compile_expr(right)?;
        let opcode = match op {
            BinaryOp::Add => OpCode::Add,
            BinaryOp::Sub => OpCode::Sub,
            BinaryOp::Mul => OpCode::Mul,
            BinaryOp::Div => OpCode::Div,
            BinaryOp::Mod => OpCode::Mod,
            BinaryOp::Greater => OpCode::Greater,
            BinaryOp::GreaterEqual => OpCode::GreaterEqual,
            BinaryOp::Less => OpCode::Less,
            BinaryOp::LessEqual => OpCode::LessEqual,
            BinaryOp::Equal => OpCode::Equal,
            BinaryOp::NotEqual => OpCode::NotEqual,
        };
        self.emit_op(opcode);
        Ok(())
    }

    fn compile_logical(&mut self, left: &Expr, op: &LogicalOp, right: &Expr) -> Result<(), RuntimeError> {
        self.compile_expr(left)?;
        match op {
            LogicalOp::And => {
                // Short-circuit: if left is falsy, skip right and return left
                let end_jump = self.emit_jump(OpCode::PopJumpIfFalse);
                self.emit_op(OpCode::Pop); // discard left (truthy)
                self.compile_expr(right)?;
                self.patch_jump(end_jump);
            }
            LogicalOp::Or => {
                // Short-circuit: if left is truthy, skip right and return left
                self.emit_op(OpCode::Dup);
                let end_jump = self.emit_jump(OpCode::JumpIfTrue);
                self.emit_op(OpCode::Pop); // discard left (falsy)
                self.compile_expr(right)?;
                self.patch_jump(end_jump);
            }
        }
        Ok(())
    }

    fn compile_ternary(&mut self, condition: &Expr, then_branch: &Expr, else_branch: &Expr) -> Result<(), RuntimeError> {
        self.compile_expr(condition)?;
        let else_jump = self.emit_jump(OpCode::JumpIfFalse);
        self.compile_expr(then_branch)?;
        let end_jump = self.emit_jump(OpCode::Jump);
        self.patch_jump(else_jump);
        self.compile_expr(else_branch)?;
        self.patch_jump(end_jump);
        Ok(())
    }

    fn compile_call(&mut self, callee: &Expr, arguments: &[Expr]) -> Result<(), RuntimeError> {
        self.compile_expr(callee)?;
        for arg in arguments {
            self.compile_expr(arg)?;
        }
        emit_u8(&mut self.function.code, OpCode::Call, arguments.len() as u8);
        Ok(())
    }

    fn compile_lambda(&mut self, params: &[String], body: &[Stmt]) -> Result<(), RuntimeError> {
        let proto = self.compile_nested_function(None, params, body)?;
        let upvalues = proto.upvalues.clone();
        let proto_idx = self.add_function_proto_constant(proto);
        self.emit_closure(proto_idx, &upvalues);
        Ok(())
    }

    /// Compile a nested function (lambda or function declaration) and return its proto.
    fn compile_nested_function(&mut self, name: Option<String>, params: &[String], body: &[Stmt]) -> Result<FunctionProto, RuntimeError> {
        let mut child = Compiler::new(name);
        // Build merged enclosing locals: our locals + our enclosing chain
        let mut all_locals = self.locals.clone();
        if let Some(ref enc) = self.enclosing_locals {
            all_locals.extend(enc.iter().cloned());
        }
        child.enclosing_locals = Some(all_locals);
        child.parent_locals_count = self.locals.len();
        child.enclosing_upvalues = Some(Rc::clone(&self.upvalues));
        child.grandparent_upvalues = self.enclosing_upvalues.clone();

        // Add params as locals without StoreLocal — they're already on stack from Call
        for param in params {
            child.locals.push(Local {
                name: param.clone(),
                depth: 1, // params are at scope depth 1 (function body)
                is_captured: false,
                is_const: false,
            });
        }
        child.function.arity = params.len() as u8;

        // Compile the body
        child.scope_depth = 1;
        for stmt in body {
            child.compile_stmt(stmt)?;
        }
        // Implicit return nil
        child.emit_op(OpCode::LoadNil);
        child.emit_op(OpCode::Return);

        // Mark captured locals in the parent compiler
        for uv in child.upvalues.borrow().iter() {
            if uv.is_local {
                if let Some(local) = self.locals.get_mut(uv.index as usize) {
                    local.is_captured = true;
                }
            }
        }

        let child_upvalues = child.upvalues.borrow();
        child.function.upvalue_count = child_upvalues.len() as u8;
        child.function.upvalues = child_upvalues.iter().map(|uv| (uv.is_local, uv.index)).collect();
        drop(child_upvalues);
        Ok(child.function)
    }

    // ========================================================================
    // Scope management
    // ========================================================================

    fn begin_scope(&mut self) {
        self.scope_depth += 1;
    }

    fn end_scope(&mut self) {
        self.scope_depth -= 1;
        // Pop locals that are going out of scope (except captured ones)
        let mut pop_count = 0u8;
        while let Some(local) = self.locals.last() {
            if local.depth > self.scope_depth {
                if local.is_captured {
                    // Don't pop — the upvalue still needs it on the stack
                    // The VM will close the upvalue when the function returns
                    self.locals.pop();
                } else {
                    self.locals.pop();
                    pop_count += 1;
                }
            } else {
                break;
            }
        }
        for _ in 0..pop_count {
            self.emit_op(OpCode::Pop);
        }
    }

    // ========================================================================
    // Variable resolution
    // ========================================================================

    /// Find a local variable by name, returning its slot index.
    fn resolve_local(&self, name: &str) -> Option<u8> {
        for (i, local) in self.locals.iter().enumerate().rev() {
            if local.name == name && local.depth > 0 {
                return Some(i as u8);
            }
        }
        None
    }

    /// Try to resolve a variable as an upvalue from enclosing functions.
    /// Returns the upvalue index in this function's upvalues list.
    fn resolve_upvalue(&mut self, name: &str) -> Option<u8> {
        // Check if already captured
        for (j, uv) in self.upvalues.borrow().iter().enumerate() {
            if uv.name == name {
                return Some(j as u8);
            }
        }
        // Check all enclosing locals (merged chain: parent, grandparent, ...)
        if let Some(ref enclosing_locals) = self.enclosing_locals {
            for (i, local) in enclosing_locals.iter().enumerate().rev() {
                if local.name == name && local.depth > 0 {
                    if i < self.parent_locals_count {
                        // Found in direct parent's locals → capture as upvalue
                        let idx = self.upvalues.borrow().len() as u8;
                        self.upvalues.borrow_mut().push(Upvalue {
                            index: i as u8, is_local: true,
                            is_const: local.is_const,
                            name: name.to_string()
                        });
                        return Some(idx);
                    } else {
                        // Found deeper than parent. Need to create upvalue chain.
                        if let Some(ref enc_upvalues) = self.enclosing_upvalues {
                            // Check if parent already has this upvalue
                            let parent_uv_idx = {
                                let enc = enc_upvalues.borrow();
                                enc.iter().position(|uv| uv.name == name).map(|p| p as u8)
                            };
                            let parent_idx = match parent_uv_idx {
                                Some(idx) => idx,
                                None => {
                                    // Add upvalue to parent's list.
                                    // The parent sees this variable at a certain index in its own
                                    // enclosing_locals. That index is (i - self.parent_locals_count)
                                    // in the merged list, which corresponds to the same variable
                                    // in the parent's enclosing_locals.
                                    let enc_idx = enc_upvalues.borrow().len() as u8;

                                    // If we have grandparent_upvalues, the parent's upvalue
                                    // should be transitive (is_local=false), and we need to
                                    // ensure grandparent has it too.
                                    if let Some(ref gp_upvalues) = self.grandparent_upvalues {
                                        // Ensure grandparent has the upvalue first
                                        let gp_idx = {
                                            let gp = gp_upvalues.borrow();
                                            gp.iter().position(|uv| uv.name == name).map(|p| p as u8)
                                        };
                                        let gp_idx = match gp_idx {
                                            Some(idx) => idx,
                                            None => {
                                                let idx = gp_upvalues.borrow().len() as u8;
                                                // Grandparent captures this as a local
                                                // (it's directly in the grandparent's enclosing scope)
                                                gp_upvalues.borrow_mut().push(Upvalue {
                                                    index: (i - self.parent_locals_count) as u8,
                                                    is_local: true,
                                                    is_const: false, // can't easily resolve const-ness at this depth
                                                    name: name.to_string()
                                                });
                                                idx
                                            }
                                        };
                                        // Parent's upvalue is transitive through grandparent
                                        let gp_uv_is_const = gp_upvalues.borrow()[gp_idx as usize].is_const;
                                        enc_upvalues.borrow_mut().push(Upvalue {
                                            index: gp_idx,
                                            is_local: false,
                                            is_const: gp_uv_is_const,
                                            name: name.to_string()
                                        });
                                    } else {
                                        // No grandparent — parent captures directly as local
                                        enc_upvalues.borrow_mut().push(Upvalue {
                                            index: (i - self.parent_locals_count) as u8,
                                            is_local: true,
                                            is_const: local.is_const,
                                            name: name.to_string()
                                        });
                                    }
                                    enc_idx
                                }
                            };
                            // Add transitive upvalue in self pointing to parent's
                            let parent_uv_is_const = enc_upvalues.borrow()[parent_idx as usize].is_const;
                            let idx = self.upvalues.borrow().len() as u8;
                            self.upvalues.borrow_mut().push(Upvalue {
                                index: parent_idx, is_local: false,
                                is_const: parent_uv_is_const,
                                name: name.to_string()
                            });
                            return Some(idx);
                        }
                    }
                }
            }
        }
        // Check parent's upvalues (transitive closure over 2 levels)
        if let Some(ref enclosing_upvalues) = self.enclosing_upvalues {
            for uv in enclosing_upvalues.borrow().iter().rev() {
                if uv.name == name {
                    let idx = self.upvalues.borrow().len() as u8;
                    self.upvalues.borrow_mut().push(Upvalue {
                        index: uv.index, is_local: false,
                        is_const: uv.is_const,
                        name: name.to_string()
                    });
                    return Some(idx);
                }
            }
        }
        None
    }

    // ========================================================================
    // Bytecode emission helpers
    // ========================================================================

    fn emit_op(&mut self, op: OpCode) {
        emit_op(&mut self.function.code, op);
    }

    fn emit_jump(&mut self, op: OpCode) -> usize {
        let loc = self.function.code.len();
        emit_i16(&mut self.function.code, op, 0x7FFF); // placeholder
        loc
    }

    fn emit_jump_placeholder(&mut self) -> usize {
        let loc = self.function.code.len();
        emit_i16(&mut self.function.code, OpCode::Jump, 0x7FFF);
        loc
    }

    fn emit_loop_jump(&mut self, target: usize) {
        let offset = target as isize - self.function.code.len() as isize;
        emit_i16(&mut self.function.code, OpCode::Jump, offset as i16);
    }

    fn emit_forin_next(&mut self, items_slot: u8, idx_slot: u8) -> usize {
        let loc = self.function.code.len();
        let code = &mut self.function.code;
        code.push(OpCode::ForInNext as u8);
        code.push(items_slot);
        code.push(idx_slot);
        code.push(0xFF); // placeholder offset low
        code.push(0x7F); // placeholder offset high
        loc
    }

    fn emit_closure(&mut self, proto_idx: u16, upvalues: &[(bool, u8)]) {
        let code = &mut self.function.code;
        code.push(OpCode::Closure as u8);
        code.push((proto_idx & 0xFF) as u8);
        code.push(((proto_idx >> 8) & 0xFF) as u8);
        let upvalue_count = upvalues.len() as u8;
        code.push(upvalue_count);
        for &(is_local, index) in upvalues {
            code.push(if is_local { 1u8 } else { 0u8 });
            code.push(index);
        }
    }

    fn patch_jump(&mut self, jump_loc: usize) {
        let offset = (self.function.code.len() - jump_loc) as i16;
        let code = &mut self.function.code;
        code[jump_loc + 1] = (offset & 0xFF) as u8;
        code[jump_loc + 2] = ((offset >> 8) & 0xFF) as u8;
    }

    fn patch_jump_to(&mut self, jump_loc: usize, target: usize) {
        let offset = target as isize - jump_loc as isize;
        let code = &mut self.function.code;
        code[jump_loc + 1] = (offset & 0xFF) as u8;
        code[jump_loc + 2] = ((offset >> 8) & 0xFF) as u8;
    }

    /// Patch ForInNext's exit offset (at jump_loc + 3, +4)
    fn patch_forin_jump(&mut self, forin_loc: usize, target: usize) {
        let offset = target as isize - forin_loc as isize;
        let code = &mut self.function.code;
        code[forin_loc + 3] = (offset & 0xFF) as u8;
        code[forin_loc + 4] = ((offset >> 8) & 0xFF) as u8;
    }

    fn add_constant(&mut self, val: Value) -> u16 {
        self.function.add_constant(val)
    }

    // ========================================================================
    // Constant pool helpers
    // ========================================================================

    fn add_string_constant(&mut self, s: &str) -> u16 {
        self.function.add_constant(Value::String(s.to_string()))
    }

    fn add_function_proto_constant(&mut self, proto: FunctionProto) -> u16 {
        let idx = self.function.protos.len() as u16;
        self.function.protos.push(Rc::new(proto));
        // Store proto index as a sentinel value in constants
        self.function.constants.push(Value::Number(f64::from_bits(idx as u64 | 0x_F000_0000_0000_0000)));
        idx
    }
}

// ============================================================================
// Helpers
// ============================================================================

fn assign_op_to_compound(op: &AssignOp) -> CompoundOp {
    match op {
        AssignOp::PlusEqual => CompoundOp::PlusEqual,
        AssignOp::MinusEqual => CompoundOp::MinusEqual,
        AssignOp::StarEqual => CompoundOp::StarEqual,
        AssignOp::SlashEqual => CompoundOp::SlashEqual,
        AssignOp::PercentEqual => CompoundOp::PercentEqual,
        AssignOp::Equal => unreachable!(),
    }
}