use crate::ast::*;
use crate::error::CompileError;
use crate::opcodes::OpCode;
use crate::token::Span;
use crate::ui_decl::{UiDeclaration, UiElement};
use crate::validator::VType;
use crate::vm::ScriptVM;

/// Where a named variable lives in the VM
#[derive(Debug, Clone, Copy)]
enum VarLoc {
    Local(u16, VType),
    Param(u16),
    StateScalar(u16, VType),
    InputBuffer(u8),
    OutputBuffer(u8),
    StateArray(u16, VType), // VType is the element type
    SampleSlot(u8),
    BuiltinSampleRate,
    BuiltinBufferSize,
}

struct Compiler {
    code: Vec<u8>,
    constants_f32: Vec<f32>,
    constants_i32: Vec<i32>,
    vars: Vec<(String, VarLoc)>,
    next_local: u16,
    scope_stack: Vec<u16>, // local count at scope entry
    draw_context: bool,     // true when compiling a draw {} block
}

impl Compiler {
    fn new() -> Self {
        Self {
            code: Vec::new(),
            constants_f32: Vec::new(),
            constants_i32: Vec::new(),
            vars: Vec::new(),
            next_local: 0,
            scope_stack: Vec::new(),
            draw_context: false,
        }
    }

    fn emit(&mut self, op: OpCode) {
        self.code.push(op as u8);
    }

    fn emit_u8(&mut self, v: u8) {
        self.code.push(v);
    }

    fn emit_u16(&mut self, v: u16) {
        self.code.extend_from_slice(&v.to_le_bytes());
    }

    fn emit_u32(&mut self, v: u32) {
        self.code.extend_from_slice(&v.to_le_bytes());
    }

    /// Returns index into constants_f32
    fn add_const_f32(&mut self, v: f32) -> u16 {
        // Reuse existing constant if possible
        for (i, &c) in self.constants_f32.iter().enumerate() {
            if c.to_bits() == v.to_bits() {
                return i as u16;
            }
        }
        let idx = self.constants_f32.len() as u16;
        self.constants_f32.push(v);
        idx
    }

    /// Returns index into constants_i32
    fn add_const_i32(&mut self, v: i32) -> u16 {
        for (i, &c) in self.constants_i32.iter().enumerate() {
            if c == v {
                return i as u16;
            }
        }
        let idx = self.constants_i32.len() as u16;
        self.constants_i32.push(v);
        idx
    }

    fn push_scope(&mut self) {
        self.scope_stack.push(self.next_local);
    }

    fn pop_scope(&mut self) {
        let prev = self.scope_stack.pop().unwrap();
        // Remove variables defined in this scope
        self.vars.retain(|(_, loc)| {
            if let VarLoc::Local(idx, _) = loc {
                *idx < prev
            } else {
                true
            }
        });
        self.next_local = prev;
    }

    fn alloc_local(&mut self, name: String, ty: VType) -> u16 {
        let idx = self.next_local;
        self.next_local += 1;
        self.vars.push((name, VarLoc::Local(idx, ty)));
        idx
    }

    fn lookup(&self, name: &str) -> Option<VarLoc> {
        self.vars.iter().rev().find(|(n, _)| n == name).map(|(_, l)| *l)
    }

    /// Emit a placeholder u32 and return the offset where it was written
    fn emit_jump_placeholder(&mut self, op: OpCode) -> usize {
        self.emit(op);
        let pos = self.code.len();
        self.emit_u32(0);
        pos
    }

    /// Patch a previously emitted u32 placeholder
    fn patch_jump(&mut self, placeholder_pos: usize) {
        let target = self.code.len() as u32;
        let bytes = target.to_le_bytes();
        self.code[placeholder_pos] = bytes[0];
        self.code[placeholder_pos + 1] = bytes[1];
        self.code[placeholder_pos + 2] = bytes[2];
        self.code[placeholder_pos + 3] = bytes[3];
    }

    fn compile_script(&mut self, script: &Script) -> Result<(), CompileError> {
        // Register built-in variables
        self.vars.push(("sample_rate".into(), VarLoc::BuiltinSampleRate));
        self.vars.push(("buffer_size".into(), VarLoc::BuiltinBufferSize));

        // Register inputs
        for (i, input) in script.inputs.iter().enumerate() {
            match input.signal {
                SignalKind::Audio | SignalKind::Cv => {
                    self.vars.push((input.name.clone(), VarLoc::InputBuffer(i as u8)));
                }
                SignalKind::Midi => {}
            }
        }

        // Register outputs
        for (i, output) in script.outputs.iter().enumerate() {
            match output.signal {
                SignalKind::Audio | SignalKind::Cv => {
                    self.vars.push((output.name.clone(), VarLoc::OutputBuffer(i as u8)));
                }
                SignalKind::Midi => {}
            }
        }

        self.register_params_and_state(script, true);

        // Compile process block
        for stmt in &script.process {
            self.compile_stmt(stmt)?;
        }

        self.emit(OpCode::Halt);
        Ok(())
    }

    /// Compile the draw block into separate bytecode (for the DrawVM)
    fn compile_draw(&mut self, script: &Script) -> Result<(), CompileError> {
        self.draw_context = true;
        self.register_params_and_state(script, false);

        // Compile draw block
        if let Some(draw) = &script.draw {
            for stmt in draw {
                self.compile_stmt(stmt)?;
            }
        }

        self.emit(OpCode::Halt);
        Ok(())
    }

    /// Register params and state variables. If `include_samples` is true, also registers sample slots.
    fn register_params_and_state(&mut self, script: &Script, include_samples: bool) {
        for (i, param) in script.params.iter().enumerate() {
            self.vars.push((param.name.clone(), VarLoc::Param(i as u16)));
        }

        let mut scalar_idx: u16 = 0;
        let mut array_idx: u16 = 0;
        let mut sample_idx: u8 = 0;
        for state in &script.state {
            match &state.ty {
                StateType::F32 => {
                    self.vars.push((state.name.clone(), VarLoc::StateScalar(scalar_idx, VType::F32)));
                    scalar_idx += 1;
                }
                StateType::Int => {
                    self.vars.push((state.name.clone(), VarLoc::StateScalar(scalar_idx, VType::Int)));
                    scalar_idx += 1;
                }
                StateType::Bool => {
                    self.vars.push((state.name.clone(), VarLoc::StateScalar(scalar_idx, VType::Bool)));
                    scalar_idx += 1;
                }
                StateType::ArrayF32(_) => {
                    self.vars.push((state.name.clone(), VarLoc::StateArray(array_idx, VType::F32)));
                    array_idx += 1;
                }
                StateType::ArrayInt(_) => {
                    self.vars.push((state.name.clone(), VarLoc::StateArray(array_idx, VType::Int)));
                    array_idx += 1;
                }
                StateType::Sample if include_samples => {
                    self.vars.push((state.name.clone(), VarLoc::SampleSlot(sample_idx)));
                    sample_idx += 1;
                }
                StateType::Sample => {}
            }
        }
    }

    fn compile_stmt(&mut self, stmt: &Stmt) -> Result<(), CompileError> {
        match stmt {
            Stmt::Let { name, init, .. } => {
                let ty = self.infer_type(init)?;
                self.compile_expr(init)?;
                let _idx = self.alloc_local(name.clone(), ty);
                self.emit(OpCode::StoreLocal);
                self.emit_u16(_idx);
            }
            Stmt::Assign { target, value, span } => {
                match target {
                    LValue::Ident(name, _) => {
                        let loc = self.lookup(name).ok_or_else(|| {
                            CompileError::new(format!("Undefined variable: {}", name), *span)
                        })?;
                        self.compile_expr(value)?;
                        match loc {
                            VarLoc::Local(idx, _) => {
                                self.emit(OpCode::StoreLocal);
                                self.emit_u16(idx);
                            }
                            VarLoc::StateScalar(idx, _) => {
                                self.emit(OpCode::StoreState);
                                self.emit_u16(idx);
                            }
                            VarLoc::Param(idx) if self.draw_context => {
                                self.emit(OpCode::StoreParam);
                                self.emit_u16(idx);
                            }
                            _ => {
                                return Err(CompileError::new(
                                    format!("Cannot assign to {}", name), *span,
                                ));
                            }
                        }
                    }
                    LValue::Index(name, idx_expr, s) => {
                        let loc = self.lookup(name).ok_or_else(|| {
                            CompileError::new(format!("Undefined variable: {}", name), *s)
                        })?;
                        match loc {
                            VarLoc::OutputBuffer(port) => {
                                // StoreOutput: pops value then index
                                self.compile_expr(idx_expr)?;
                                self.compile_expr(value)?;
                                self.emit(OpCode::StoreOutput);
                                self.emit_u8(port);
                            }
                            VarLoc::StateArray(arr_id, _) => {
                                // StoreStateArray: pops value then index
                                self.compile_expr(idx_expr)?;
                                self.compile_expr(value)?;
                                self.emit(OpCode::StoreStateArray);
                                self.emit_u16(arr_id);
                            }
                            _ => {
                                return Err(CompileError::new(
                                    format!("Cannot index-assign to {}", name), *s,
                                ));
                            }
                        }
                    }
                }
            }
            Stmt::If { cond, then_block, else_block, .. } => {
                self.compile_expr(cond)?;
                if let Some(else_b) = else_block {
                    // JumpIfFalse -> else
                    let else_jump = self.emit_jump_placeholder(OpCode::JumpIfFalse);
                    self.push_scope();
                    self.compile_block(then_block)?;
                    self.pop_scope();
                    // Jump -> end (skip else)
                    let end_jump = self.emit_jump_placeholder(OpCode::Jump);
                    self.patch_jump(else_jump);
                    self.push_scope();
                    self.compile_block(else_b)?;
                    self.pop_scope();
                    self.patch_jump(end_jump);
                } else {
                    let end_jump = self.emit_jump_placeholder(OpCode::JumpIfFalse);
                    self.push_scope();
                    self.compile_block(then_block)?;
                    self.pop_scope();
                    self.patch_jump(end_jump);
                }
            }
            Stmt::For { var, end, body, span: _ } => {
                // Allocate loop variable as local
                self.push_scope();
                let loop_var = self.alloc_local(var.clone(), VType::Int);

                // Initialize loop var to 0
                let zero_idx = self.add_const_i32(0);
                self.emit(OpCode::PushI32);
                self.emit_u16(zero_idx);
                self.emit(OpCode::StoreLocal);
                self.emit_u16(loop_var);

                // Loop start: check condition (i < end)
                let loop_start = self.code.len();
                self.emit(OpCode::LoadLocal);
                self.emit_u16(loop_var);
                self.compile_expr(end)?;
                self.emit(OpCode::LtI);

                let exit_jump = self.emit_jump_placeholder(OpCode::JumpIfFalse);

                // Body
                self.compile_block(body)?;

                // Increment loop var
                self.emit(OpCode::LoadLocal);
                self.emit_u16(loop_var);
                let one_idx = self.add_const_i32(1);
                self.emit(OpCode::PushI32);
                self.emit_u16(one_idx);
                self.emit(OpCode::AddI);
                self.emit(OpCode::StoreLocal);
                self.emit_u16(loop_var);

                // Jump back to loop start
                self.emit(OpCode::Jump);
                self.emit_u32(loop_start as u32);

                // Patch exit
                self.patch_jump(exit_jump);
                self.pop_scope();
            }
            Stmt::ExprStmt(expr) => {
                let is_void = self.is_void_call(expr);
                self.compile_expr(expr)?;
                if !is_void {
                    self.emit(OpCode::Pop);
                }
            }
        }
        Ok(())
    }

    fn compile_block(&mut self, block: &[Stmt]) -> Result<(), CompileError> {
        for stmt in block {
            self.compile_stmt(stmt)?;
        }
        Ok(())
    }

    fn compile_expr(&mut self, expr: &Expr) -> Result<(), CompileError> {
        match expr {
            Expr::FloatLit(v, _) => {
                let idx = self.add_const_f32(*v);
                self.emit(OpCode::PushF32);
                self.emit_u16(idx);
            }
            Expr::IntLit(v, _) => {
                let idx = self.add_const_i32(*v);
                self.emit(OpCode::PushI32);
                self.emit_u16(idx);
            }
            Expr::BoolLit(v, _) => {
                self.emit(OpCode::PushBool);
                self.emit_u8(if *v { 1 } else { 0 });
            }
            Expr::Ident(name, span) => {
                let loc = self.lookup(name).ok_or_else(|| {
                    CompileError::new(format!("Undefined variable: {}", name), *span)
                })?;
                match loc {
                    VarLoc::Local(idx, _) => {
                        self.emit(OpCode::LoadLocal);
                        self.emit_u16(idx);
                    }
                    VarLoc::Param(idx) => {
                        self.emit(OpCode::LoadParam);
                        self.emit_u16(idx);
                    }
                    VarLoc::StateScalar(idx, _) => {
                        self.emit(OpCode::LoadState);
                        self.emit_u16(idx);
                    }
                    VarLoc::BuiltinSampleRate => {
                        self.emit(OpCode::LoadSampleRate);
                    }
                    VarLoc::BuiltinBufferSize => {
                        self.emit(OpCode::LoadBufferSize);
                    }
                    // Arrays/buffers/samples used bare (for len(), etc.) — handled by call codegen
                    _ => {}
                }
            }
            Expr::BinOp(left, op, right, _span) => {
                let lt = self.infer_type(left)?;
                let rt = self.infer_type(right)?;
                self.compile_expr(left)?;
                self.compile_expr(right)?;

                match op {
                    BinOp::Add => {
                        if lt == VType::F32 || rt == VType::F32 {
                            self.emit(OpCode::AddF);
                        } else {
                            self.emit(OpCode::AddI);
                        }
                    }
                    BinOp::Sub => {
                        if lt == VType::F32 || rt == VType::F32 {
                            self.emit(OpCode::SubF);
                        } else {
                            self.emit(OpCode::SubI);
                        }
                    }
                    BinOp::Mul => {
                        if lt == VType::F32 || rt == VType::F32 {
                            self.emit(OpCode::MulF);
                        } else {
                            self.emit(OpCode::MulI);
                        }
                    }
                    BinOp::Div => {
                        if lt == VType::F32 || rt == VType::F32 {
                            self.emit(OpCode::DivF);
                        } else {
                            self.emit(OpCode::DivI);
                        }
                    }
                    BinOp::Mod => {
                        if lt == VType::F32 || rt == VType::F32 {
                            self.emit(OpCode::ModF);
                        } else {
                            self.emit(OpCode::ModI);
                        }
                    }
                    BinOp::Eq => {
                        if lt == VType::F32 || rt == VType::F32 {
                            self.emit(OpCode::EqF);
                        } else if lt == VType::Int || rt == VType::Int {
                            self.emit(OpCode::EqI);
                        } else {
                            // bool comparison: treat as int
                            self.emit(OpCode::EqI);
                        }
                    }
                    BinOp::Ne => {
                        if lt == VType::F32 || rt == VType::F32 {
                            self.emit(OpCode::NeF);
                        } else {
                            self.emit(OpCode::NeI);
                        }
                    }
                    BinOp::Lt => {
                        if lt == VType::F32 || rt == VType::F32 {
                            self.emit(OpCode::LtF);
                        } else {
                            self.emit(OpCode::LtI);
                        }
                    }
                    BinOp::Gt => {
                        if lt == VType::F32 || rt == VType::F32 {
                            self.emit(OpCode::GtF);
                        } else {
                            self.emit(OpCode::GtI);
                        }
                    }
                    BinOp::Le => {
                        if lt == VType::F32 || rt == VType::F32 {
                            self.emit(OpCode::LeF);
                        } else {
                            self.emit(OpCode::LeI);
                        }
                    }
                    BinOp::Ge => {
                        if lt == VType::F32 || rt == VType::F32 {
                            self.emit(OpCode::GeF);
                        } else {
                            self.emit(OpCode::GeI);
                        }
                    }
                    BinOp::And => self.emit(OpCode::And),
                    BinOp::Or => self.emit(OpCode::Or),
                }
            }
            Expr::UnaryOp(op, inner, _) => {
                let ty = self.infer_type(inner)?;
                self.compile_expr(inner)?;
                match op {
                    UnaryOp::Neg => {
                        if ty == VType::F32 {
                            self.emit(OpCode::NegF);
                        } else {
                            self.emit(OpCode::NegI);
                        }
                    }
                    UnaryOp::Not => self.emit(OpCode::Not),
                }
            }
            Expr::Cast(kind, inner, _) => {
                self.compile_expr(inner)?;
                match kind {
                    CastKind::ToInt => self.emit(OpCode::F32ToI32),
                    CastKind::ToFloat => self.emit(OpCode::I32ToF32),
                }
            }
            Expr::Index(base, idx, span) => {
                // base must be an Ident referencing an array/buffer
                if let Expr::Ident(name, _) = base.as_ref() {
                    let loc = self.lookup(name).ok_or_else(|| {
                        CompileError::new(format!("Undefined variable: {}", name), *span)
                    })?;
                    match loc {
                        VarLoc::InputBuffer(port) => {
                            self.compile_expr(idx)?;
                            self.emit(OpCode::LoadInput);
                            self.emit_u8(port);
                        }
                        VarLoc::OutputBuffer(port) => {
                            self.compile_expr(idx)?;
                            self.emit(OpCode::LoadInput);
                            // Reading from output buffer — use same port but from outputs
                            // Actually outputs aren't readable in the VM. This would be
                            // an error in practice, but the validator should catch it.
                            // For now, treat as input read (will read zeros).
                            self.emit_u8(port);
                        }
                        VarLoc::StateArray(arr_id, _) => {
                            self.compile_expr(idx)?;
                            self.emit(OpCode::LoadStateArray);
                            self.emit_u16(arr_id);
                        }
                        _ => {
                            return Err(CompileError::new(
                                format!("Cannot index variable: {}", name), *span,
                            ));
                        }
                    }
                } else {
                    return Err(CompileError::new("Index base must be an identifier", *span));
                }
            }
            Expr::Call(name, args, span) => {
                self.compile_call(name, args, *span)?;
            }
        }
        Ok(())
    }

    /// Returns true if the expression is a call to a void function (no return value).
    fn is_void_call(&self, expr: &Expr) -> bool {
        if let Expr::Call(name, _, _) = expr {
            matches!(name.as_str(),
                "fill_circle" | "stroke_circle" | "stroke_arc" |
                "line" | "fill_rect" | "stroke_rect"
            )
        } else {
            false
        }
    }

    fn compile_call(&mut self, name: &str, args: &[Expr], span: Span) -> Result<(), CompileError> {
        match name {
            // 1-arg math → push arg, emit opcode
            "sin" => { self.compile_expr(&args[0])?; self.emit(OpCode::Sin); }
            "cos" => { self.compile_expr(&args[0])?; self.emit(OpCode::Cos); }
            "tan" => { self.compile_expr(&args[0])?; self.emit(OpCode::Tan); }
            "asin" => { self.compile_expr(&args[0])?; self.emit(OpCode::Asin); }
            "acos" => { self.compile_expr(&args[0])?; self.emit(OpCode::Acos); }
            "atan" => { self.compile_expr(&args[0])?; self.emit(OpCode::Atan); }
            "exp" => { self.compile_expr(&args[0])?; self.emit(OpCode::Exp); }
            "log" => { self.compile_expr(&args[0])?; self.emit(OpCode::Log); }
            "log2" => { self.compile_expr(&args[0])?; self.emit(OpCode::Log2); }
            "sqrt" => { self.compile_expr(&args[0])?; self.emit(OpCode::Sqrt); }
            "floor" => { self.compile_expr(&args[0])?; self.emit(OpCode::Floor); }
            "ceil" => { self.compile_expr(&args[0])?; self.emit(OpCode::Ceil); }
            "round" => { self.compile_expr(&args[0])?; self.emit(OpCode::Round); }
            "trunc" => { self.compile_expr(&args[0])?; self.emit(OpCode::Trunc); }
            "fract" => { self.compile_expr(&args[0])?; self.emit(OpCode::Fract); }
            "abs" => { self.compile_expr(&args[0])?; self.emit(OpCode::Abs); }
            "sign" => { self.compile_expr(&args[0])?; self.emit(OpCode::Sign); }

            // 2-arg math
            "atan2" => {
                self.compile_expr(&args[0])?;
                self.compile_expr(&args[1])?;
                self.emit(OpCode::Atan2);
            }
            "pow" => {
                self.compile_expr(&args[0])?;
                self.compile_expr(&args[1])?;
                self.emit(OpCode::Pow);
            }
            "min" => {
                self.compile_expr(&args[0])?;
                self.compile_expr(&args[1])?;
                self.emit(OpCode::Min);
            }
            "max" => {
                self.compile_expr(&args[0])?;
                self.compile_expr(&args[1])?;
                self.emit(OpCode::Max);
            }

            // 3-arg math
            "clamp" => {
                self.compile_expr(&args[0])?;
                self.compile_expr(&args[1])?;
                self.compile_expr(&args[2])?;
                self.emit(OpCode::Clamp);
            }
            "mix" => {
                self.compile_expr(&args[0])?;
                self.compile_expr(&args[1])?;
                self.compile_expr(&args[2])?;
                self.emit(OpCode::Mix);
            }
            "smoothstep" => {
                self.compile_expr(&args[0])?;
                self.compile_expr(&args[1])?;
                self.compile_expr(&args[2])?;
                self.emit(OpCode::Smoothstep);
            }

            // cv_or(value, default) — if value is NaN, use default
            "cv_or" => {
                // Compile: push value, check IsNan, if true use default else keep value
                // Strategy: push value, dup-like via local, IsNan, branch
                // Simpler: push value, push value again, IsNan, JumpIfFalse skip, Pop, push default, skip:
                // But we don't have Dup. Use a temp local instead.
                let temp = self.next_local;
                self.next_local += 1;
                self.compile_expr(&args[0])?;
                // Store to temp
                self.emit(OpCode::StoreLocal);
                self.emit_u16(temp);
                // Load and check NaN
                self.emit(OpCode::LoadLocal);
                self.emit_u16(temp);
                self.emit(OpCode::IsNan);
                let skip_default = self.emit_jump_placeholder(OpCode::JumpIfFalse);
                // NaN path: use default
                self.compile_expr(&args[1])?;
                let skip_end = self.emit_jump_placeholder(OpCode::Jump);
                // Not NaN path: use original value
                self.patch_jump(skip_default);
                self.emit(OpCode::LoadLocal);
                self.emit_u16(temp);
                self.patch_jump(skip_end);
                self.next_local -= 1; // release temp
            }

            // len(array) -> int
            "len" => {
                // Arg must be an ident referencing a state array or input/output buffer
                if let Expr::Ident(arr_name, s) = &args[0] {
                    let loc = self.lookup(arr_name).ok_or_else(|| {
                        CompileError::new(format!("Undefined variable: {}", arr_name), *s)
                    })?;
                    match loc {
                        VarLoc::StateArray(arr_id, _) => {
                            self.emit(OpCode::ArrayLen);
                            self.emit_u16(arr_id);
                        }
                        VarLoc::InputBuffer(_) | VarLoc::OutputBuffer(_) => {
                            // Buffer length is buffer_size (for CV) or buffer_size*2 (for audio)
                            // We emit LoadBufferSize — scripts use buffer_size for iteration
                            self.emit(OpCode::LoadBufferSize);
                        }
                        _ => {
                            return Err(CompileError::new("len() argument must be an array", span));
                        }
                    }
                } else {
                    return Err(CompileError::new("len() argument must be an identifier", span));
                }
            }

            // sample_len(sample) -> int
            "sample_len" => {
                if let Expr::Ident(sname, s) = &args[0] {
                    let loc = self.lookup(sname).ok_or_else(|| {
                        CompileError::new(format!("Undefined: {}", sname), *s)
                    })?;
                    if let VarLoc::SampleSlot(slot) = loc {
                        self.emit(OpCode::SampleLen);
                        self.emit_u8(slot);
                    } else {
                        return Err(CompileError::new("sample_len() requires a sample", span));
                    }
                } else {
                    return Err(CompileError::new("sample_len() requires an identifier", span));
                }
            }

            // sample_read(sample, index) -> f32
            "sample_read" => {
                if let Expr::Ident(sname, s) = &args[0] {
                    let loc = self.lookup(sname).ok_or_else(|| {
                        CompileError::new(format!("Undefined: {}", sname), *s)
                    })?;
                    if let VarLoc::SampleSlot(slot) = loc {
                        self.compile_expr(&args[1])?;
                        self.emit(OpCode::SampleRead);
                        self.emit_u8(slot);
                    } else {
                        return Err(CompileError::new("sample_read() requires a sample", span));
                    }
                } else {
                    return Err(CompileError::new("sample_read() requires an identifier", span));
                }
            }

            // sample_rate_of(sample) -> int
            "sample_rate_of" => {
                if let Expr::Ident(sname, s) = &args[0] {
                    let loc = self.lookup(sname).ok_or_else(|| {
                        CompileError::new(format!("Undefined: {}", sname), *s)
                    })?;
                    if let VarLoc::SampleSlot(slot) = loc {
                        self.emit(OpCode::SampleRateOf);
                        self.emit_u8(slot);
                    } else {
                        return Err(CompileError::new("sample_rate_of() requires a sample", span));
                    }
                } else {
                    return Err(CompileError::new("sample_rate_of() requires an identifier", span));
                }
            }

            // Draw builtins (only valid in draw context)
            "fill_circle" | "stroke_circle" | "stroke_arc" | "line" |
            "fill_rect" | "stroke_rect" |
            "mouse_x" | "mouse_y" | "mouse_down" if self.draw_context => {
                match name {
                    "fill_circle" => {
                        // fill_circle(cx, cy, r, color)
                        for arg in args { self.compile_expr(arg)?; }
                        self.emit(OpCode::DrawFillCircle);
                    }
                    "stroke_circle" => {
                        // stroke_circle(cx, cy, r, color, width)
                        for arg in args { self.compile_expr(arg)?; }
                        self.emit(OpCode::DrawStrokeCircle);
                    }
                    "stroke_arc" => {
                        // stroke_arc(cx, cy, r, start_deg, end_deg, color, width)
                        for arg in args { self.compile_expr(arg)?; }
                        self.emit(OpCode::DrawStrokeArc);
                    }
                    "line" => {
                        // line(x1, y1, x2, y2, color, width)
                        for arg in args { self.compile_expr(arg)?; }
                        self.emit(OpCode::DrawLine);
                    }
                    "fill_rect" => {
                        // fill_rect(x, y, w, h, color)
                        for arg in args { self.compile_expr(arg)?; }
                        self.emit(OpCode::DrawFillRect);
                    }
                    "stroke_rect" => {
                        // stroke_rect(x, y, w, h, color, width)
                        for arg in args { self.compile_expr(arg)?; }
                        self.emit(OpCode::DrawStrokeRect);
                    }
                    "mouse_x" => { self.emit(OpCode::MouseX); }
                    "mouse_y" => { self.emit(OpCode::MouseY); }
                    "mouse_down" => { self.emit(OpCode::MouseDown); }
                    _ => unreachable!(),
                }
            }

            _ => {
                return Err(CompileError::new(format!("Unknown function: {}", name), span));
            }
        }
        Ok(())
    }

    /// Infer the type of an expression (mirrors validator logic, needed for selecting typed opcodes)
    fn infer_type(&self, expr: &Expr) -> Result<VType, CompileError> {
        match expr {
            Expr::FloatLit(_, _) => Ok(VType::F32),
            Expr::IntLit(_, _) => Ok(VType::Int),
            Expr::BoolLit(_, _) => Ok(VType::Bool),
            Expr::Ident(name, span) => {
                let loc = self.lookup(name).ok_or_else(|| {
                    CompileError::new(format!("Undefined variable: {}", name), *span)
                })?;
                match loc {
                    VarLoc::Local(_, ty) => Ok(ty),
                    VarLoc::Param(_) => Ok(VType::F32),
                    VarLoc::StateScalar(_, ty) => Ok(ty),
                    VarLoc::InputBuffer(_) => Ok(VType::ArrayF32),
                    VarLoc::OutputBuffer(_) => Ok(VType::ArrayF32),
                    VarLoc::StateArray(_, elem_ty) => {
                        if elem_ty == VType::Int { Ok(VType::ArrayInt) } else { Ok(VType::ArrayF32) }
                    }
                    VarLoc::SampleSlot(_) => Ok(VType::Sample),
                    VarLoc::BuiltinSampleRate => Ok(VType::Int),
                    VarLoc::BuiltinBufferSize => Ok(VType::Int),
                }
            }
            Expr::BinOp(left, op, right, _) => {
                let lt = self.infer_type(left)?;
                let rt = self.infer_type(right)?;
                match op {
                    BinOp::And | BinOp::Or | BinOp::Eq | BinOp::Ne |
                    BinOp::Lt | BinOp::Gt | BinOp::Le | BinOp::Ge => Ok(VType::Bool),
                    _ => {
                        if lt == VType::F32 || rt == VType::F32 {
                            Ok(VType::F32)
                        } else {
                            Ok(VType::Int)
                        }
                    }
                }
            }
            Expr::UnaryOp(op, inner, _) => {
                match op {
                    UnaryOp::Neg => self.infer_type(inner),
                    UnaryOp::Not => Ok(VType::Bool),
                }
            }
            Expr::Cast(kind, _, _) => match kind {
                CastKind::ToInt => Ok(VType::Int),
                CastKind::ToFloat => Ok(VType::F32),
            },
            Expr::Index(base, _, _) => {
                let base_ty = self.infer_type(base)?;
                match base_ty {
                    VType::ArrayF32 => Ok(VType::F32),
                    VType::ArrayInt => Ok(VType::Int),
                    _ => Ok(VType::F32), // fallback
                }
            }
            Expr::Call(name, _, _) => {
                match name.as_str() {
                    "len" | "sample_len" | "sample_rate_of" => Ok(VType::Int),
                    "isnan" => Ok(VType::Bool),
                    _ => Ok(VType::F32), // all math functions return f32
                }
            }
        }
    }
}

/// Compile a validated AST into bytecode VM and UI declaration
pub fn compile(script: &Script) -> Result<(ScriptVM, UiDeclaration, Option<crate::vm::DrawVM>), CompileError> {
    let mut compiler = Compiler::new();
    compiler.compile_script(script)?;

    // Collect state layout info
    let mut num_state_scalars = 0usize;
    let mut state_array_sizes = Vec::new();
    let mut num_sample_slots = 0usize;

    for state in &script.state {
        match &state.ty {
            StateType::F32 | StateType::Int | StateType::Bool => {
                num_state_scalars += 1;
            }
            StateType::ArrayF32(sz) => state_array_sizes.push(*sz),
            StateType::ArrayInt(sz) => state_array_sizes.push(*sz),
            StateType::Sample => num_sample_slots += 1,
        }
    }

    let param_defaults: Vec<f32> = script.params.iter().map(|p| p.default).collect();

    let vm = ScriptVM::new(
        compiler.code,
        compiler.constants_f32,
        compiler.constants_i32,
        script.params.len(),
        &param_defaults,
        num_state_scalars,
        &state_array_sizes,
        num_sample_slots,
    );

    // Build UI declaration
    let ui_decl = if let Some(elements) = &script.ui {
        UiDeclaration { elements: elements.clone() }
    } else {
        // Auto-generate: sample pickers first, then all params
        let mut elements = Vec::new();
        for state in &script.state {
            if state.ty == StateType::Sample {
                elements.push(UiElement::Sample(state.name.clone()));
            }
        }
        for param in &script.params {
            elements.push(UiElement::Param(param.name.clone()));
        }
        UiDeclaration { elements }
    };

    // Compile draw block if present
    let draw_vm = if script.draw.is_some() {
        let mut draw_compiler = Compiler::new();
        draw_compiler.compile_draw(script)?;
        Some(crate::vm::DrawVM::new(
            draw_compiler.code,
            draw_compiler.constants_f32,
            draw_compiler.constants_i32,
            script.params.len(),
            &param_defaults,
            num_state_scalars,
            &state_array_sizes,
        ))
    } else {
        None
    };

    Ok((vm, ui_decl, draw_vm))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::lexer::Lexer;
    use crate::parser::Parser;
    use crate::validator;

    fn compile_source(src: &str) -> Result<(ScriptVM, UiDeclaration), CompileError> {
        let mut lexer = Lexer::new(src);
        let tokens = lexer.tokenize()?;
        let mut parser = Parser::new(&tokens);
        let script = parser.parse()?;
        let validated = validator::validate(&script)?;
        let (vm, ui, _draw_vm) = compile(validated)?;
        Ok((vm, ui))
    }

    #[test]
    fn test_passthrough() {
        let src = r#"
            name "Pass"
            category effect
            inputs { audio_in: audio }
            outputs { audio_out: audio }
            process {
                for i in 0..buffer_size {
                    audio_out[i] = audio_in[i];
                }
            }
        "#;
        let (mut vm, _) = compile_source(src).unwrap();
        let input = vec![1.0f32, 2.0, 3.0, 4.0];
        let mut output = vec![0.0f32; 4];
        let inputs: Vec<&[f32]> = vec![&input];
        let mut out_slice: Vec<&mut [f32]> = vec![&mut output];
        vm.execute(&inputs, &mut out_slice, 44100, 4).unwrap();
        assert_eq!(output, vec![1.0, 2.0, 3.0, 4.0]);
    }

    #[test]
    fn test_gain() {
        let src = r#"
            name "Gain"
            category effect
            inputs { audio_in: audio }
            outputs { audio_out: audio }
            params { gain: 0.5 [0.0, 1.0] "" }
            process {
                for i in 0..buffer_size {
                    audio_out[i] = audio_in[i] * gain;
                }
            }
        "#;
        let (mut vm, _) = compile_source(src).unwrap();
        let input = vec![1.0f32, 2.0, 3.0, 4.0];
        let mut output = vec![0.0f32; 4];
        let inputs: Vec<&[f32]> = vec![&input];
        let mut out_slice: Vec<&mut [f32]> = vec![&mut output];
        vm.execute(&inputs, &mut out_slice, 44100, 4).unwrap();
        assert_eq!(output, vec![0.5, 1.0, 1.5, 2.0]);
    }

    #[test]
    fn test_state_array() {
        let src = r#"
            name "Delay"
            category effect
            inputs { audio_in: audio }
            outputs { audio_out: audio }
            state { buf: [8]f32 }
            process {
                for i in 0..buffer_size {
                    audio_out[i] = buf[i];
                    buf[i] = audio_in[i];
                }
            }
        "#;
        let (mut vm, _) = compile_source(src).unwrap();

        // First call: output should be zeros (state initialized to 0), state gets input
        let input = vec![10.0f32, 20.0, 30.0, 40.0];
        let mut output = vec![0.0f32; 4];
        {
            let inputs: Vec<&[f32]> = vec![&input];
            let mut out_slice: Vec<&mut [f32]> = vec![&mut output];
            vm.execute(&inputs, &mut out_slice, 44100, 4).unwrap();
        }
        assert_eq!(output, vec![0.0, 0.0, 0.0, 0.0]);

        // Second call: output should be previous input
        let input2 = vec![50.0f32, 60.0, 70.0, 80.0];
        let mut output2 = vec![0.0f32; 4];
        {
            let inputs: Vec<&[f32]> = vec![&input2];
            let mut out_slice: Vec<&mut [f32]> = vec![&mut output2];
            vm.execute(&inputs, &mut out_slice, 44100, 4).unwrap();
        }
        assert_eq!(output2, vec![10.0, 20.0, 30.0, 40.0]);
    }

    #[test]
    fn test_if_else() {
        let src = r#"
            name "Gate"
            category effect
            inputs { audio_in: audio }
            outputs { audio_out: audio }
            params { threshold: 0.5 [0.0, 1.0] "" }
            process {
                for i in 0..buffer_size {
                    if audio_in[i] >= threshold {
                        audio_out[i] = audio_in[i];
                    } else {
                        audio_out[i] = 0.0;
                    }
                }
            }
        "#;
        let (mut vm, _) = compile_source(src).unwrap();
        let input = vec![0.2f32, 0.8, 0.1, 0.9];
        let mut output = vec![0.0f32; 4];
        let inputs: Vec<&[f32]> = vec![&input];
        let mut out_slice: Vec<&mut [f32]> = vec![&mut output];
        vm.execute(&inputs, &mut out_slice, 44100, 4).unwrap();
        assert_eq!(output, vec![0.0, 0.8, 0.0, 0.9]);
    }

    #[test]
    fn test_auto_ui() {
        let src = r#"
            name "Test"
            category utility
            params { gain: 1.0 [0.0, 2.0] "dB" }
            state { clip: sample }
            outputs { out: audio }
            process {}
        "#;
        let (_, ui) = compile_source(src).unwrap();
        // Auto-generated: sample first, then params
        assert_eq!(ui.elements.len(), 2);
        assert!(matches!(&ui.elements[0], UiElement::Sample(n) if n == "clip"));
        assert!(matches!(&ui.elements[1], UiElement::Param(n) if n == "gain"));
    }

    #[test]
    fn test_draw_block() {
        let src = r#"
            name "Knob"
            category utility
            params { volume: 0.75 [0.0, 1.0] "" }
            outputs { out: audio }
            ui { canvas [80, 80] }
            draw {
                let cx = 40.0;
                let cy = 40.0;
                fill_circle(cx, cy, 35.0, 0x333333FF);
                let angle = volume * 270.0 - 135.0;
                stroke_arc(cx, cy, 30.0, -135.0, angle, 0x4488FFFF, 3.0);
            }
            process {
                for i in 0..buffer_size {
                    out[i] = 0.0;
                }
            }
        "#;
        let mut lexer = Lexer::new(src);
        let tokens = lexer.tokenize().unwrap();
        let mut parser = Parser::new(&tokens);
        let script = parser.parse().unwrap();
        let validated = validator::validate(&script).unwrap();
        let (_vm, _ui, draw_vm) = compile(validated).unwrap();

        // Draw VM should exist
        let mut dvm = draw_vm.expect("draw_vm should be Some");
        assert!(dvm.has_bytecode());

        // Execute should succeed without stack errors
        dvm.execute().unwrap();

        // Should have produced draw commands
        assert_eq!(dvm.draw_commands.len(), 2); // fill_circle + stroke_arc

    }
}