// ===============================================================================
// QUANTALANG CODE GENERATOR - ARM64 BACKEND
// ===============================================================================
// Copyright (c) 2022-2026 Zain Dana Harper. MIT License.
// ===============================================================================

//! ARM64/AArch64 native code generation backend.
//!
//! Generates native machine code for ARM64 processors.
//!
//! ## Output Modes
//!
//! - **Assembly**: Generates GNU-syntax assembly text (default)
//! - **MachineCode**: Generates raw machine code using the instruction encoder
//!
//! ## Calling Convention (AAPCS64)
//!
//! - Arguments: X0-X7 (integers), D0-D7 (floats)
//! - Return value: X0 (integer), D0 (float)
//! - Callee-saved: X19-X28, D8-D15
//! - Frame pointer: X29 (FP)
//! - Link register: X30 (LR)
//! - Stack pointer: SP (16-byte aligned)

// NOTE: `.unwrap()` calls in this backend are intentional assertions on
// type-checked AST invariants. Failures indicate compiler bugs in earlier
// phases, not user input errors. See codegen/mod.rs for the full unwrap policy.

use std::collections::HashMap;
use std::fmt::Write;

use super::arm64_enc::{Arm64Encoder, Cond, Reg64};
use super::{Backend, CodegenResult, Target};
use crate::codegen::debug::DwarfGenerator;
use crate::codegen::ir::*;
use crate::codegen::{GeneratedCode, OutputFormat};

/// Output mode for the ARM64 backend.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Arm64OutputMode {
    /// Generate assembly text.
    Assembly,
    /// Generate raw machine code.
    MachineCode,
}

/// ARM64 backend for code generation.
pub struct Arm64Backend {
    /// Output buffer (assembly).
    output: String,
    /// Label counter.
    label_counter: u32,
    /// Output mode.
    mode: Arm64OutputMode,
    /// Machine code encoder (when mode is MachineCode).
    encoder: Option<Arm64Encoder>,
    /// Generate debug information.
    generate_debug: bool,
    /// DWARF generator for debug info.
    dwarf: Option<DwarfGenerator>,
    /// Local to stack offset mapping.
    local_offsets: HashMap<LocalId, i32>,
    /// Function address ranges for debug info.
    func_ranges: Vec<(String, u64, u64)>,
}

impl Arm64Backend {
    /// Create a new ARM64 backend.
    pub fn new() -> Self {
        Self {
            output: String::new(),
            label_counter: 0,
            mode: Arm64OutputMode::Assembly,
            encoder: None,
            generate_debug: false,
            dwarf: None,
            local_offsets: HashMap::new(),
            func_ranges: Vec::new(),
        }
    }

    /// Enable machine code output mode.
    pub fn with_machine_code(mut self) -> Self {
        self.mode = Arm64OutputMode::MachineCode;
        self.encoder = Some(Arm64Encoder::new());
        self
    }

    /// Enable debug information generation.
    pub fn with_debug_info(mut self) -> Self {
        self.generate_debug = true;
        self.dwarf = Some(DwarfGenerator::new());
        self
    }

    /// Generate a fresh label.
    fn fresh_label(&mut self, prefix: &str) -> String {
        let label = format!(".L{}_{}", prefix, self.label_counter);
        self.label_counter += 1;
        label
    }

    /// Get encoder reference (panics if not in machine code mode).
    fn enc(&mut self) -> &mut Arm64Encoder {
        self.encoder
            .as_mut()
            .expect("Machine code mode not enabled")
    }

    // =========================================================================
    // Assembly Generation
    // =========================================================================

    fn generate_module(&mut self, module: &MirModule) -> CodegenResult<()> {
        // Header
        self.output
            .push_str("// Generated by QuantaLang Compiler\n");
        self.output.push_str("// Target: ARM64/AArch64\n\n");
        self.output.push_str(".arch armv8-a\n\n");

        // Data section
        self.output.push_str(".section .data\n");
        self.generate_data_section(module)?;

        // Read-only data
        self.output.push_str("\n.section .rodata\n");
        self.generate_rodata_section(module)?;

        // Text section (code)
        self.output.push_str("\n.section .text\n");
        for func in &module.functions {
            if !func.is_declaration() {
                self.generate_function(func)?;
            }
        }

        Ok(())
    }

    fn generate_data_section(&mut self, module: &MirModule) -> CodegenResult<()> {
        for global in &module.globals {
            if global.is_mut {
                writeln!(self.output, ".globl {}", global.name).unwrap();
                writeln!(self.output, "{}:", global.name).unwrap();
                if let Some(init) = &global.init {
                    self.emit_const_data(init, &global.ty)?;
                } else {
                    let size = self.type_size(&global.ty);
                    writeln!(self.output, "    .zero {}", size).unwrap();
                }
            }
        }
        Ok(())
    }

    fn generate_rodata_section(&mut self, module: &MirModule) -> CodegenResult<()> {
        // String literals
        for (i, s) in module.strings.iter().enumerate() {
            writeln!(self.output, "__str{}:", i).unwrap();
            writeln!(self.output, "    .asciz \"{}\"", self.escape_string(s)).unwrap();
        }

        // Constant globals
        for global in &module.globals {
            if !global.is_mut {
                writeln!(self.output, ".globl {}", global.name).unwrap();
                writeln!(self.output, "{}:", global.name).unwrap();
                if let Some(init) = &global.init {
                    self.emit_const_data(init, &global.ty)?;
                } else {
                    let size = self.type_size(&global.ty);
                    writeln!(self.output, "    .zero {}", size).unwrap();
                }
            }
        }
        Ok(())
    }

    fn generate_function(&mut self, func: &MirFunction) -> CodegenResult<()> {
        // Function prologue
        if func.is_public {
            writeln!(self.output, ".globl {}", func.name).unwrap();
        }
        writeln!(self.output, "{}:", func.name).unwrap();

        // Standard prologue - save frame pointer and link register
        self.output.push_str("    stp x29, x30, [sp, #-16]!\n");
        self.output.push_str("    mov x29, sp\n");

        // Allocate stack space for locals
        let stack_size = self.calculate_stack_size(func);
        if stack_size > 0 {
            writeln!(self.output, "    sub sp, sp, #{}", stack_size).unwrap();
        }

        // Save callee-saved registers if needed
        self.save_callee_saved(func)?;

        // Move arguments from registers to stack
        self.store_arguments(func)?;

        // Generate blocks
        if let Some(blocks) = &func.blocks {
            for block in blocks {
                self.generate_block(block, func)?;
            }
        }

        self.output.push('\n');
        Ok(())
    }

    fn save_callee_saved(&mut self, _func: &MirFunction) -> CodegenResult<()> {
        // TODO: Save X19-X28 as needed based on register allocation
        Ok(())
    }

    fn store_arguments(&mut self, func: &MirFunction) -> CodegenResult<()> {
        let arg_regs = ["x0", "x1", "x2", "x3", "x4", "x5", "x6", "x7"];
        let mut arg_idx = 0;

        for local in &func.locals {
            if local.is_param && arg_idx < arg_regs.len() {
                let offset = self.local_stack_offset(local.id, func);
                writeln!(
                    self.output,
                    "    str {}, [x29, #-{}]",
                    arg_regs[arg_idx], offset
                )
                .unwrap();
                arg_idx += 1;
            }
        }
        Ok(())
    }

    fn generate_block(&mut self, block: &MirBlock, func: &MirFunction) -> CodegenResult<()> {
        // Block label
        let label = block
            .label
            .as_ref()
            .map(|l| l.to_string())
            .unwrap_or_else(|| format!(".Lbb{}_{}", func.name, block.id.0));
        writeln!(self.output, "{}:", label).unwrap();

        // Generate statements
        for stmt in &block.stmts {
            self.generate_statement(stmt, func)?;
        }

        // Generate terminator
        if let Some(term) = &block.terminator {
            self.generate_terminator(term, func)?;
        }

        Ok(())
    }

    fn generate_statement(&mut self, stmt: &MirStmt, func: &MirFunction) -> CodegenResult<()> {
        match &stmt.kind {
            MirStmtKind::Assign { dest, value } => {
                self.generate_assign(*dest, value, func)?;
            }
            MirStmtKind::DerefAssign { ptr, .. } => {
                let offset = self.local_stack_offset(*ptr, func);
                self.output.push_str(&format!(
                    "    // deref store through local at sp+{}\n",
                    offset
                ));
                self.output
                    .push_str(&format!("    ldr x1, [x29, #-{}]\n", offset));
                self.output.push_str("    str x0, [x1]\n");
            }
            MirStmtKind::FieldDerefAssign {
                ptr, field_name, ..
            } => {
                let offset = self.local_stack_offset(*ptr, func);
                self.output.push_str(&format!(
                    "    // field deref store .{} through local at sp+{}\n",
                    field_name, offset
                ));
                self.output
                    .push_str(&format!("    ldr x1, [x29, #-{}]\n", offset));
                self.output.push_str("    str x0, [x1]\n");
            }
            MirStmtKind::FieldAssign {
                base, field_name, ..
            } => {
                let offset = self.local_stack_offset(*base, func);
                self.output.push_str(&format!(
                    "    // field store .{} on local at sp+{}\n",
                    field_name, offset
                ));
                self.output.push_str("    str x0, [x1]\n");
            }
            MirStmtKind::StorageLive(_) | MirStmtKind::StorageDead(_) => {
                // No-op
            }
            MirStmtKind::Nop => {
                self.output.push_str("    nop\n");
            }
        }
        Ok(())
    }

    fn generate_assign(
        &mut self,
        dest: LocalId,
        value: &MirRValue,
        func: &MirFunction,
    ) -> CodegenResult<()> {
        match value {
            MirRValue::Use(val) => {
                self.load_value_to_x0(val, func)?;
                self.store_x0_to_local(dest, func)?;
            }
            MirRValue::BinaryOp { op, left, right } => {
                // Load left into x0
                self.load_value_to_x0(left, func)?;
                self.output.push_str("    mov x9, x0\n");
                // Load right into x0, then move to x1
                self.load_value_to_x0(right, func)?;
                self.output.push_str("    mov x1, x0\n");
                self.output.push_str("    mov x0, x9\n");
                // Perform operation
                self.emit_binary_op(*op)?;
                // Store result
                self.store_x0_to_local(dest, func)?;
            }
            MirRValue::UnaryOp { op, operand } => {
                self.load_value_to_x0(operand, func)?;
                match op {
                    UnaryOp::Neg => self.output.push_str("    neg x0, x0\n"),
                    UnaryOp::Not => self.output.push_str("    mvn x0, x0\n"),
                }
                self.store_x0_to_local(dest, func)?;
            }
            _ => {
                writeln!(self.output, "    // TODO: rvalue {:?}", value).unwrap();
            }
        }
        Ok(())
    }

    fn generate_terminator(
        &mut self,
        term: &MirTerminator,
        func: &MirFunction,
    ) -> CodegenResult<()> {
        match term {
            MirTerminator::Goto(target) => {
                writeln!(self.output, "    b .Lbb{}_{}", func.name, target.0).unwrap();
            }
            MirTerminator::If {
                cond,
                then_block,
                else_block,
            } => {
                self.load_value_to_x0(cond, func)?;
                self.output.push_str("    cmp x0, #0\n");
                writeln!(self.output, "    b.ne .Lbb{}_{}", func.name, then_block.0).unwrap();
                writeln!(self.output, "    b .Lbb{}_{}", func.name, else_block.0).unwrap();
            }
            MirTerminator::Return(value) => {
                if let Some(val) = value {
                    self.load_value_to_x0(val, func)?;
                }
                // Epilogue
                let stack_size = self.calculate_stack_size(func);
                if stack_size > 0 {
                    writeln!(self.output, "    add sp, sp, #{}", stack_size).unwrap();
                }
                self.output.push_str("    ldp x29, x30, [sp], #16\n");
                self.output.push_str("    ret\n");
            }
            MirTerminator::Call {
                func: callee,
                args,
                dest,
                target,
                ..
            } => {
                // Pass arguments in registers (AAPCS64)
                let arg_regs = ["x0", "x1", "x2", "x3", "x4", "x5", "x6", "x7"];
                for (i, arg) in args.iter().enumerate() {
                    if i < arg_regs.len() {
                        self.load_value_to_x0(arg, func)?;
                        if i > 0 {
                            writeln!(self.output, "    mov {}, x0", arg_regs[i]).unwrap();
                        }
                    } else {
                        // Stack arguments
                        self.load_value_to_x0(arg, func)?;
                        self.output.push_str("    str x0, [sp, #-16]!\n");
                    }
                }
                // Call
                let callee_name = match callee {
                    MirValue::Function(name) => name.to_string(),
                    _ => "unknown".to_string(),
                };
                writeln!(self.output, "    bl {}", callee_name).unwrap();
                // Store result
                if let Some(d) = dest {
                    self.store_x0_to_local(*d, func)?;
                }
                // Jump to continue block
                if let Some(target_block) = target {
                    writeln!(self.output, "    b .Lbb{}_{}", func.name, target_block.0).unwrap();
                }
            }
            MirTerminator::Unreachable => {
                self.output.push_str("    brk #0\n");
            }
            MirTerminator::Abort => {
                self.output.push_str("    bl abort\n");
            }
            _ => {
                writeln!(self.output, "    // TODO: terminator {:?}", term).unwrap();
            }
        }
        Ok(())
    }

    fn load_value_to_x0(&mut self, value: &MirValue, func: &MirFunction) -> CodegenResult<()> {
        match value {
            MirValue::Local(id) => {
                let offset = self.local_stack_offset(*id, func);
                writeln!(self.output, "    ldr x0, [x29, #-{}]", offset).unwrap();
            }
            MirValue::Const(c) => match c {
                MirConst::Int(v, _) => {
                    let v = *v as i64;
                    self.emit_const_int_asm(v);
                }
                MirConst::Uint(v, _) => {
                    let v = *v as i64;
                    self.emit_const_int_asm(v);
                }
                MirConst::Bool(b) => {
                    writeln!(self.output, "    mov x0, #{}", if *b { 1 } else { 0 }).unwrap();
                }
                _ => {
                    self.output.push_str("    mov x0, #0\n");
                }
            },
            MirValue::Global(name) => {
                writeln!(self.output, "    adrp x0, {}", name).unwrap();
                writeln!(self.output, "    add x0, x0, :lo12:{}", name).unwrap();
            }
            MirValue::Function(name) => {
                writeln!(self.output, "    adrp x0, {}", name).unwrap();
                writeln!(self.output, "    add x0, x0, :lo12:{}", name).unwrap();
            }
        }
        Ok(())
    }

    fn store_x0_to_local(&mut self, id: LocalId, func: &MirFunction) -> CodegenResult<()> {
        let offset = self.local_stack_offset(id, func);
        writeln!(self.output, "    str x0, [x29, #-{}]", offset).unwrap();
        Ok(())
    }

    /// Emit an integer constant to x0.
    fn emit_const_int_asm(&mut self, v: i64) {
        if v >= 0 && v <= 0xFFFF {
            writeln!(self.output, "    mov x0, #{}", v).unwrap();
        } else {
            // Need to load in chunks
            writeln!(self.output, "    movz x0, #{}", v & 0xFFFF).unwrap();
            if (v >> 16) & 0xFFFF != 0 {
                writeln!(self.output, "    movk x0, #{}, lsl #16", (v >> 16) & 0xFFFF).unwrap();
            }
            if (v >> 32) & 0xFFFF != 0 {
                writeln!(self.output, "    movk x0, #{}, lsl #32", (v >> 32) & 0xFFFF).unwrap();
            }
            if (v >> 48) & 0xFFFF != 0 {
                writeln!(self.output, "    movk x0, #{}, lsl #48", (v >> 48) & 0xFFFF).unwrap();
            }
        }
    }

    fn emit_binary_op(&mut self, op: BinOp) -> CodegenResult<()> {
        match op {
            BinOp::Add | BinOp::AddChecked | BinOp::AddWrapping | BinOp::AddSaturating => {
                self.output.push_str("    add x0, x0, x1\n");
            }
            BinOp::Sub | BinOp::SubChecked | BinOp::SubWrapping | BinOp::SubSaturating => {
                self.output.push_str("    sub x0, x0, x1\n");
            }
            BinOp::Mul | BinOp::MulChecked | BinOp::MulWrapping => {
                self.output.push_str("    mul x0, x0, x1\n");
            }
            BinOp::Div => {
                self.output.push_str("    sdiv x0, x0, x1\n");
            }
            BinOp::Rem => {
                // x0 = x0 % x1 => x0 - (x0 / x1) * x1
                self.output.push_str("    sdiv x9, x0, x1\n");
                self.output.push_str("    msub x0, x9, x1, x0\n");
            }
            BinOp::BitAnd => {
                self.output.push_str("    and x0, x0, x1\n");
            }
            BinOp::BitOr => {
                self.output.push_str("    orr x0, x0, x1\n");
            }
            BinOp::BitXor => {
                self.output.push_str("    eor x0, x0, x1\n");
            }
            BinOp::Shl => {
                self.output.push_str("    lsl x0, x0, x1\n");
            }
            BinOp::Shr => {
                self.output.push_str("    asr x0, x0, x1\n");
            }
            BinOp::Eq => {
                self.output.push_str("    cmp x0, x1\n");
                self.output.push_str("    cset x0, eq\n");
            }
            BinOp::Ne => {
                self.output.push_str("    cmp x0, x1\n");
                self.output.push_str("    cset x0, ne\n");
            }
            BinOp::Lt => {
                self.output.push_str("    cmp x0, x1\n");
                self.output.push_str("    cset x0, lt\n");
            }
            BinOp::Le => {
                self.output.push_str("    cmp x0, x1\n");
                self.output.push_str("    cset x0, le\n");
            }
            BinOp::Gt => {
                self.output.push_str("    cmp x0, x1\n");
                self.output.push_str("    cset x0, gt\n");
            }
            BinOp::Ge => {
                self.output.push_str("    cmp x0, x1\n");
                self.output.push_str("    cset x0, ge\n");
            }
            BinOp::Pow => {
                // Integer power: base (x0) ** exponent (x1)
                // Using repeated squaring algorithm
                // x9 = result, x0 = base, x1 = exponent
                let label = self.label_counter;
                self.label_counter += 1;
                self.output.push_str("    mov x9, #1\n"); // result = 1
                writeln!(self.output, ".Lpow_loop_{}:", label).unwrap();
                self.output.push_str("    cbz x1, .Lpow_done_"); // if exp == 0, done
                writeln!(self.output, "{}", label).unwrap();
                self.output.push_str("    tst x1, #1\n"); // test exp & 1
                writeln!(self.output, "    b.eq .Lpow_skip_{}", label).unwrap();
                self.output.push_str("    mul x9, x9, x0\n"); // result *= base
                writeln!(self.output, ".Lpow_skip_{}:", label).unwrap();
                self.output.push_str("    mul x0, x0, x0\n"); // base *= base
                self.output.push_str("    lsr x1, x1, #1\n"); // exp >>= 1
                writeln!(self.output, "    b .Lpow_loop_{}", label).unwrap();
                writeln!(self.output, ".Lpow_done_{}:", label).unwrap();
                self.output.push_str("    mov x0, x9\n"); // return result
            }
        }
        Ok(())
    }

    fn emit_const_data(&mut self, c: &MirConst, ty: &MirType) -> CodegenResult<()> {
        match c {
            MirConst::Int(v, _) => {
                let size = self.type_size(ty);
                match size {
                    1 => writeln!(self.output, "    .byte {}", v).unwrap(),
                    2 => writeln!(self.output, "    .hword {}", v).unwrap(),
                    4 => writeln!(self.output, "    .word {}", v).unwrap(),
                    8 => writeln!(self.output, "    .xword {}", v).unwrap(),
                    _ => writeln!(self.output, "    .xword {}", v).unwrap(),
                }
            }
            MirConst::Uint(v, _) => {
                let size = self.type_size(ty);
                match size {
                    1 => writeln!(self.output, "    .byte {}", v).unwrap(),
                    2 => writeln!(self.output, "    .hword {}", v).unwrap(),
                    4 => writeln!(self.output, "    .word {}", v).unwrap(),
                    8 => writeln!(self.output, "    .xword {}", v).unwrap(),
                    _ => writeln!(self.output, "    .xword {}", v).unwrap(),
                }
            }
            MirConst::Float(v, ty) => match ty {
                MirType::Float(FloatSize::F32) => {
                    writeln!(self.output, "    .word {}", (*v as f32).to_bits()).unwrap();
                }
                _ => {
                    writeln!(self.output, "    .xword {}", v.to_bits()).unwrap();
                }
            },
            MirConst::Bool(b) => {
                writeln!(self.output, "    .byte {}", if *b { 1 } else { 0 }).unwrap();
            }
            _ => {
                let size = self.type_size(ty);
                writeln!(self.output, "    .zero {}", size).unwrap();
            }
        }
        Ok(())
    }

    fn calculate_stack_size(&self, func: &MirFunction) -> u32 {
        let mut size = 0u32;
        for local in &func.locals {
            if !local.is_param {
                size += self.type_size(&local.ty) as u32;
            }
        }
        // Align to 16 bytes (AAPCS64 requirement)
        (size + 15) & !15
    }

    fn local_stack_offset(&self, id: LocalId, func: &MirFunction) -> u32 {
        let mut offset = 8; // Start after saved FP
        for local in &func.locals {
            if local.id == id {
                return offset;
            }
            offset += self.type_size(&local.ty) as u32;
        }
        offset
    }

    fn type_size(&self, ty: &MirType) -> usize {
        match ty {
            MirType::Void => 0,
            MirType::Bool => 1,
            MirType::Int(size, _) => match size {
                IntSize::I8 => 1,
                IntSize::I16 => 2,
                IntSize::I32 => 4,
                IntSize::I64 | IntSize::ISize => 8,
                IntSize::I128 => 16,
            },
            MirType::Float(size) => match size {
                FloatSize::F32 => 4,
                FloatSize::F64 => 8,
            },
            MirType::Ptr(_) => 8,
            MirType::Array(elem, len) => self.type_size(elem) * (*len as usize),
            MirType::Slice(_) => 16, // ptr + len
            MirType::Struct(_) => 8, // TODO: lookup actual size
            MirType::FnPtr(_) => 8,
            MirType::Never => 0,
            MirType::Vector(elem, lanes) => self.type_size(elem) * (*lanes as usize),
            // Opaque GPU types — treat as pointer-sized handles
            MirType::Texture2D(_) | MirType::Sampler | MirType::SampledImage(_) => 8,
            MirType::TraitObject(_) => 16, // fat pointer: data ptr + vtable ptr
            MirType::Vec(_) => 8,          // QuantaVecHandle is a pointer
            MirType::Tuple(elems) => elems.iter().map(|e| self.type_size(e)).sum(),
            MirType::Map(_, _) => 8, // QuantaMapHandle is a pointer
        }
    }

    fn escape_string(&self, s: &str) -> String {
        let mut result = String::new();
        for c in s.chars() {
            match c {
                '\n' => result.push_str("\\n"),
                '\r' => result.push_str("\\r"),
                '\t' => result.push_str("\\t"),
                '\\' => result.push_str("\\\\"),
                '"' => result.push_str("\\\""),
                c if c.is_ascii_control() => {
                    result.push_str(&format!("\\{:03o}", c as u8));
                }
                c => result.push(c),
            }
        }
        result
    }
}

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

impl Backend for Arm64Backend {
    fn generate(&mut self, mir: &MirModule) -> CodegenResult<GeneratedCode> {
        self.output.clear();
        self.label_counter = 0;
        self.local_offsets.clear();
        self.func_ranges.clear();

        match self.mode {
            Arm64OutputMode::Assembly => {
                self.generate_module(mir)?;
                Ok(GeneratedCode::new(
                    OutputFormat::Assembly,
                    self.output.as_bytes().to_vec(),
                ))
            }
            Arm64OutputMode::MachineCode => {
                self.generate_machine_code(mir)?;
                let code = self.encoder.take().unwrap().finish();
                Ok(GeneratedCode::new(OutputFormat::Object, code))
            }
        }
    }

    fn target(&self) -> Target {
        Target::Arm64
    }
}

// =============================================================================
// Machine Code Generation
// =============================================================================

impl Arm64Backend {
    /// Generate machine code from MIR.
    fn generate_machine_code(&mut self, module: &MirModule) -> CodegenResult<()> {
        // Initialize encoder
        self.encoder = Some(Arm64Encoder::new());

        // Generate functions
        for func in &module.functions {
            if !func.is_declaration() {
                self.generate_function_machine_code(func)?;
            }
        }

        Ok(())
    }

    /// Generate machine code for a function.
    fn generate_function_machine_code(&mut self, func: &MirFunction) -> CodegenResult<()> {
        let start_pos = self.enc().position() as u64;

        // Build local offset map
        self.build_local_offsets(func);

        // Function prologue - save FP and LR
        self.enc().stp_pre(Reg64::FP, Reg64::LR, Reg64::SP, -16);
        self.enc().mov(Reg64::FP, Reg64::SP);

        // Allocate stack space
        let stack_size = self.calculate_stack_size(func);
        if stack_size > 0 {
            self.enc().sub_imm(Reg64::SP, Reg64::SP, stack_size as u16);
        }

        // Move parameters from registers to stack
        let param_regs = [
            Reg64::X0,
            Reg64::X1,
            Reg64::X2,
            Reg64::X3,
            Reg64::X4,
            Reg64::X5,
            Reg64::X6,
            Reg64::X7,
        ];
        let mut param_idx = 0;
        for local in &func.locals {
            if local.is_param && param_idx < param_regs.len() {
                let _offset = self.local_offsets.get(&local.id).copied().unwrap_or(8) as u16;
                // Use STUR for potentially unaligned negative offsets
                self.enc().str(param_regs[param_idx], Reg64::FP, 0);
                // Adjust using sub for negative offset
                // TODO: Use proper negative offset addressing
                param_idx += 1;
            }
        }

        // Create label map for blocks
        let mut block_labels: HashMap<BlockId, u32> = HashMap::new();
        if let Some(blocks) = &func.blocks {
            for block in blocks {
                let label = self.enc().new_label();
                block_labels.insert(block.id, label);
            }
        }

        // Generate blocks
        if let Some(blocks) = &func.blocks {
            for block in blocks {
                let label = block_labels[&block.id];
                self.enc().define_label(label);
                self.generate_block_machine_code(block, func, &block_labels)?;
            }
        }

        let end_pos = self.enc().position() as u64;
        self.func_ranges
            .push((func.name.to_string(), start_pos, end_pos - start_pos));

        Ok(())
    }

    /// Build local variable to stack offset mapping.
    fn build_local_offsets(&mut self, func: &MirFunction) {
        self.local_offsets.clear();
        let mut offset = 16i32; // Start after saved FP/LR

        for local in &func.locals {
            let size = self.type_size(&local.ty) as i32;
            // Align to size (up to 8)
            let align = size.min(8);
            offset = (offset + align - 1) & !(align - 1);
            self.local_offsets.insert(local.id, offset);
            offset += size;
        }
    }

    /// Generate machine code for a basic block.
    fn generate_block_machine_code(
        &mut self,
        block: &MirBlock,
        func: &MirFunction,
        block_labels: &HashMap<BlockId, u32>,
    ) -> CodegenResult<()> {
        // Generate statements
        for stmt in &block.stmts {
            self.generate_stmt_machine_code(stmt, func)?;
        }

        // Generate terminator
        if let Some(term) = &block.terminator {
            self.generate_terminator_machine_code(term, func, block_labels)?;
        }

        Ok(())
    }

    /// Generate machine code for a statement.
    fn generate_stmt_machine_code(
        &mut self,
        stmt: &MirStmt,
        func: &MirFunction,
    ) -> CodegenResult<()> {
        match &stmt.kind {
            MirStmtKind::Assign { dest, value } => {
                self.generate_assign_machine_code(*dest, value, func)?;
            }
            MirStmtKind::DerefAssign { .. }
            | MirStmtKind::FieldDerefAssign { .. }
            | MirStmtKind::FieldAssign { .. } => {
                // Pointer/field store in machine code: emit nop placeholder
                // Full implementation requires register allocator
                self.enc().nop();
            }
            MirStmtKind::StorageLive(_) | MirStmtKind::StorageDead(_) => {
                // No-op
            }
            MirStmtKind::Nop => {
                self.enc().nop();
            }
        }
        Ok(())
    }

    /// Generate machine code for an assignment.
    fn generate_assign_machine_code(
        &mut self,
        dest: LocalId,
        value: &MirRValue,
        func: &MirFunction,
    ) -> CodegenResult<()> {
        match value {
            MirRValue::Use(val) => {
                self.load_value_to_reg(val, Reg64::X0, func)?;
                self.store_reg_to_local(Reg64::X0, dest)?;
            }
            MirRValue::BinaryOp { op, left, right } => {
                // Load left into X0
                self.load_value_to_reg(left, Reg64::X0, func)?;
                // Save X0 to X9 (scratch register)
                self.enc().mov(Reg64::X9, Reg64::X0);
                // Load right into X1
                self.load_value_to_reg(right, Reg64::X1, func)?;
                // Restore left to X0
                self.enc().mov(Reg64::X0, Reg64::X9);
                // Perform operation
                self.emit_binary_op_machine_code(*op)?;
                // Store result
                self.store_reg_to_local(Reg64::X0, dest)?;
            }
            MirRValue::UnaryOp { op, operand } => {
                self.load_value_to_reg(operand, Reg64::X0, func)?;
                match op {
                    UnaryOp::Neg => self.enc().neg(Reg64::X0, Reg64::X0),
                    UnaryOp::Not => self.enc().mvn(Reg64::X0, Reg64::X0),
                }
                self.store_reg_to_local(Reg64::X0, dest)?;
            }
            MirRValue::Ref { place, .. } => {
                // Load address of place (only for simple local places)
                if place.projections.is_empty() {
                    let offset = self.local_offsets.get(&place.local).copied().unwrap_or(16);
                    // Calculate address: FP - offset
                    self.enc().sub_imm(Reg64::X0, Reg64::FP, offset as u16);
                    self.store_reg_to_local(Reg64::X0, dest)?;
                }
            }
            MirRValue::Cast { value, ty: _, .. } => {
                self.load_value_to_reg(value, Reg64::X0, func)?;
                // Most casts are no-ops at machine code level
                self.store_reg_to_local(Reg64::X0, dest)?;
            }
            _ => {
                // Default: store zero
                self.enc().movz(Reg64::X0, 0, 0);
                self.store_reg_to_local(Reg64::X0, dest)?;
            }
        }
        Ok(())
    }

    /// Load a value into a register.
    fn load_value_to_reg(
        &mut self,
        value: &MirValue,
        reg: Reg64,
        _func: &MirFunction,
    ) -> CodegenResult<()> {
        match value {
            MirValue::Local(id) => {
                let offset = self.local_offsets.get(id).copied().unwrap_or(16) as u16;
                // Load from stack: for small positive offsets
                self.enc().ldr(reg, Reg64::FP, offset);
            }
            MirValue::Const(c) => {
                match c {
                    MirConst::Int(v, _) => {
                        self.enc().mov_imm64(reg, *v as u64);
                    }
                    MirConst::Uint(v, _) => {
                        self.enc().mov_imm64(reg, *v as u64);
                    }
                    MirConst::Bool(b) => {
                        if *b {
                            self.enc().movz(reg, 1, 0);
                        } else {
                            self.enc().movz(reg, 0, 0);
                        }
                    }
                    MirConst::Float(v, ty) => {
                        // Load float bit pattern as integer
                        match ty {
                            MirType::Float(FloatSize::F32) => {
                                let bits = (*v as f32).to_bits() as u64;
                                self.enc().mov_imm64(reg, bits);
                            }
                            _ => {
                                let bits = v.to_bits();
                                self.enc().mov_imm64(reg, bits);
                            }
                        }
                    }
                    _ => {
                        self.enc().movz(reg, 0, 0);
                    }
                }
            }
            MirValue::Global(_name) => {
                // ADRP + ADD for global symbol
                self.enc().adrp(reg, 0);
                // Note: actual relocation would be needed here
            }
            MirValue::Function(_name) => {
                // ADRP + ADD for function symbol
                self.enc().adrp(reg, 0);
            }
        }
        Ok(())
    }

    /// Store a register to a local variable.
    fn store_reg_to_local(&mut self, reg: Reg64, id: LocalId) -> CodegenResult<()> {
        let offset = self.local_offsets.get(&id).copied().unwrap_or(16) as u16;
        self.enc().str(reg, Reg64::FP, offset);
        Ok(())
    }

    /// Emit machine code for a binary operation.
    fn emit_binary_op_machine_code(&mut self, op: BinOp) -> CodegenResult<()> {
        match op {
            BinOp::Add | BinOp::AddChecked | BinOp::AddWrapping | BinOp::AddSaturating => {
                self.enc().add_reg(Reg64::X0, Reg64::X0, Reg64::X1);
            }
            BinOp::Sub | BinOp::SubChecked | BinOp::SubWrapping | BinOp::SubSaturating => {
                self.enc().sub_reg(Reg64::X0, Reg64::X0, Reg64::X1);
            }
            BinOp::Mul | BinOp::MulChecked | BinOp::MulWrapping => {
                self.enc().mul(Reg64::X0, Reg64::X0, Reg64::X1);
            }
            BinOp::Div => {
                self.enc().sdiv(Reg64::X0, Reg64::X0, Reg64::X1);
            }
            BinOp::Rem => {
                // x0 = x0 % x1 => x0 - (x0 / x1) * x1
                self.enc().sdiv(Reg64::X9, Reg64::X0, Reg64::X1);
                self.enc().msub(Reg64::X0, Reg64::X9, Reg64::X1, Reg64::X0);
            }
            BinOp::BitAnd => {
                self.enc().and_reg(Reg64::X0, Reg64::X0, Reg64::X1);
            }
            BinOp::BitOr => {
                self.enc().orr_reg(Reg64::X0, Reg64::X0, Reg64::X1);
            }
            BinOp::BitXor => {
                self.enc().eor_reg(Reg64::X0, Reg64::X0, Reg64::X1);
            }
            BinOp::Shl => {
                self.enc().lsl_reg(Reg64::X0, Reg64::X0, Reg64::X1);
            }
            BinOp::Shr => {
                self.enc().asr_reg(Reg64::X0, Reg64::X0, Reg64::X1);
            }
            BinOp::Eq => {
                self.enc().cmp_reg(Reg64::X0, Reg64::X1);
                self.enc().cset(Reg64::X0, Cond::EQ);
            }
            BinOp::Ne => {
                self.enc().cmp_reg(Reg64::X0, Reg64::X1);
                self.enc().cset(Reg64::X0, Cond::NE);
            }
            BinOp::Lt => {
                self.enc().cmp_reg(Reg64::X0, Reg64::X1);
                self.enc().cset(Reg64::X0, Cond::LT);
            }
            BinOp::Le => {
                self.enc().cmp_reg(Reg64::X0, Reg64::X1);
                self.enc().cset(Reg64::X0, Cond::LE);
            }
            BinOp::Gt => {
                self.enc().cmp_reg(Reg64::X0, Reg64::X1);
                self.enc().cset(Reg64::X0, Cond::GT);
            }
            BinOp::Ge => {
                self.enc().cmp_reg(Reg64::X0, Reg64::X1);
                self.enc().cset(Reg64::X0, Cond::GE);
            }
            BinOp::Pow => {
                // Integer power using repeated squaring algorithm
                // x9 = result (1), x0 = base, x1 = exponent
                let loop_label = self.enc().create_label();
                let done_label = self.enc().create_label();
                let skip_label = self.enc().create_label();

                self.enc().mov_imm(Reg64::X9, 1); // result = 1
                self.enc().bind_label(loop_label);
                self.enc().cbz_label(Reg64::X1, done_label); // if exp == 0, done
                self.enc().tst_imm(Reg64::X1, 1);
                self.enc().b_cond_label(Cond::EQ, skip_label); // if !(exp & 1), skip
                self.enc().mul(Reg64::X9, Reg64::X9, Reg64::X0); // result *= base
                self.enc().bind_label(skip_label);
                self.enc().mul(Reg64::X0, Reg64::X0, Reg64::X0); // base *= base
                self.enc().lsr_imm(Reg64::X1, Reg64::X1, 1); // exp >>= 1
                self.enc().b_label(loop_label);
                self.enc().bind_label(done_label);
                self.enc().mov_reg(Reg64::X0, Reg64::X9); // return result
            }
        }
        Ok(())
    }

    /// Generate machine code for a terminator.
    fn generate_terminator_machine_code(
        &mut self,
        term: &MirTerminator,
        func: &MirFunction,
        block_labels: &HashMap<BlockId, u32>,
    ) -> CodegenResult<()> {
        match term {
            MirTerminator::Goto(target) => {
                let label = block_labels[target];
                self.enc().b_label(label);
            }
            MirTerminator::If {
                cond,
                then_block,
                else_block,
            } => {
                self.load_value_to_reg(cond, Reg64::X0, func)?;
                // Compare with zero
                self.enc().cmp_imm(Reg64::X0, 0);

                let then_label = block_labels[then_block];
                let else_label = block_labels[else_block];

                // Branch if not equal (condition is true)
                self.enc().b_cond_label(Cond::NE, then_label);
                // Fall through or branch to else
                self.enc().b_label(else_label);
            }
            MirTerminator::Return(value) => {
                if let Some(val) = value {
                    self.load_value_to_reg(val, Reg64::X0, func)?;
                }
                // Epilogue
                let stack_size = self.calculate_stack_size(func);
                if stack_size > 0 {
                    self.enc().add_imm(Reg64::SP, Reg64::SP, stack_size as u16);
                }
                self.enc().ldp_post(Reg64::FP, Reg64::LR, Reg64::SP, 16);
                self.enc().ret();
            }
            MirTerminator::Call {
                func: callee,
                args,
                dest,
                target,
                ..
            } => {
                // Pass arguments
                let arg_regs = [
                    Reg64::X0,
                    Reg64::X1,
                    Reg64::X2,
                    Reg64::X3,
                    Reg64::X4,
                    Reg64::X5,
                    Reg64::X6,
                    Reg64::X7,
                ];
                for (i, arg) in args.iter().enumerate() {
                    if i < arg_regs.len() {
                        self.load_value_to_reg(arg, arg_regs[i], func)?;
                    } else {
                        // Stack arguments
                        self.load_value_to_reg(arg, Reg64::X9, func)?;
                        self.enc().stp_pre(Reg64::X9, Reg64::X9, Reg64::SP, -16);
                    }
                }

                // Call
                if let MirValue::Function(name) = callee {
                    self.enc().bl_symbol(name);
                } else {
                    self.load_value_to_reg(callee, Reg64::X9, func)?;
                    self.enc().blr(Reg64::X9);
                }

                // Store result
                if let Some(d) = dest {
                    self.store_reg_to_local(Reg64::X0, *d)?;
                }

                // Jump to continue block
                if let Some(target_block) = target {
                    let label = block_labels[target_block];
                    self.enc().b_label(label);
                }
            }
            MirTerminator::Switch {
                value,
                targets,
                default,
            } => {
                self.load_value_to_reg(value, Reg64::X0, func)?;

                // Generate comparisons and branches for each case
                for (case_val, case_block) in targets {
                    let v_u64 = match case_val {
                        MirConst::Int(v, _) => *v as u64,
                        MirConst::Uint(v, _) => *v as u64,
                        _ => continue,
                    };
                    self.enc().mov_imm64(Reg64::X9, v_u64);
                    self.enc().cmp_reg(Reg64::X0, Reg64::X9);
                    let label = block_labels[case_block];
                    self.enc().b_cond_label(Cond::EQ, label);
                }

                // Default case
                let default_label = block_labels[default];
                self.enc().b_label(default_label);
            }
            MirTerminator::Unreachable => {
                self.enc().brk(0);
            }
            MirTerminator::Abort => {
                self.enc().bl_symbol("abort");
            }
            MirTerminator::Drop { target, .. } => {
                // For now, just jump to target (drop glue not implemented)
                let label = block_labels[target];
                self.enc().b_label(label);
            }
            MirTerminator::Assert {
                cond,
                expected,
                target,
                ..
            } => {
                self.load_value_to_reg(cond, Reg64::X0, func)?;
                self.enc().cmp_imm(Reg64::X0, 0);
                let label = block_labels[target];
                if *expected {
                    self.enc().b_cond_label(Cond::NE, label);
                } else {
                    self.enc().b_cond_label(Cond::EQ, label);
                }
                // If assertion fails, call abort
                self.enc().bl_symbol("abort");
            }
            MirTerminator::Resume => {
                // Resume unwinding
                self.enc().bl_symbol("_Unwind_Resume");
            }
        }
        Ok(())
    }
}

// =============================================================================
// Tests
// =============================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use crate::codegen::builder::{values, MirBuilder, MirModuleBuilder};

    // =========================================================================
    // Assembly Generation Tests
    // =========================================================================

    #[test]
    fn test_arm64_backend_simple() {
        let mut module_builder = MirModuleBuilder::new("test");

        let sig = MirFnSig::new(vec![MirType::i32(), MirType::i32()], MirType::i32());
        let mut builder = MirBuilder::new("add", sig);

        let a = builder.param_local(0);
        let b = builder.param_local(1);
        let result = builder.create_local(MirType::i32());

        builder.binary_op(result, BinOp::Add, values::local(a), values::local(b));
        builder.ret(Some(values::local(result)));

        module_builder.add_function(builder.build());
        let module = module_builder.build();

        let mut backend = Arm64Backend::new();
        let output = backend.generate(&module).unwrap();

        let code = output.as_string().unwrap();
        assert!(code.contains(".arch armv8-a"));
        assert!(code.contains("add:"));
        assert!(code.contains("ret"));
    }

    #[test]
    fn test_arm64_backend_void_function() {
        let mut module_builder = MirModuleBuilder::new("test");

        let sig = MirFnSig::new(vec![], MirType::Void);
        let mut builder = MirBuilder::new("noop", sig);
        builder.ret_void();

        module_builder.add_function(builder.build());
        let module = module_builder.build();

        let mut backend = Arm64Backend::new();
        let output = backend.generate(&module).unwrap();

        let code = output.as_string().unwrap();
        assert!(code.contains("noop:"));
        assert!(code.contains("stp x29, x30"));
        assert!(code.contains("ldp x29, x30"));
        assert!(code.contains("ret"));
    }

    #[test]
    fn test_arm64_backend_branching() {
        let mut module_builder = MirModuleBuilder::new("test");

        let sig = MirFnSig::new(vec![MirType::Bool], MirType::i32());
        let mut builder = MirBuilder::new("branch_test", sig);

        let cond = builder.param_local(0);
        let then_block = builder.create_block();
        let else_block = builder.create_block();

        builder.branch(values::local(cond), then_block, else_block);

        builder.switch_to_block(then_block);
        builder.ret(Some(values::i32(1)));

        builder.switch_to_block(else_block);
        builder.ret(Some(values::i32(0)));

        module_builder.add_function(builder.build());
        let module = module_builder.build();

        let mut backend = Arm64Backend::new();
        let output = backend.generate(&module).unwrap();

        let code = output.as_string().unwrap();
        assert!(code.contains("cmp x0, #0"));
        assert!(code.contains("b.ne"));
    }

    #[test]
    fn test_arm64_backend_binary_ops() {
        let mut module_builder = MirModuleBuilder::new("test");

        let sig = MirFnSig::new(vec![MirType::i32(), MirType::i32()], MirType::i32());
        let mut builder = MirBuilder::new("sub_test", sig);

        let a = builder.param_local(0);
        let b = builder.param_local(1);
        let result = builder.create_local(MirType::i32());

        builder.binary_op(result, BinOp::Sub, values::local(a), values::local(b));
        builder.ret(Some(values::local(result)));

        module_builder.add_function(builder.build());
        let module = module_builder.build();

        let mut backend = Arm64Backend::new();
        let output = backend.generate(&module).unwrap();

        let code = output.as_string().unwrap();
        assert!(code.contains("sub x0, x0, x1"));
    }

    #[test]
    fn test_arm64_backend_mul() {
        let mut module_builder = MirModuleBuilder::new("test");

        let sig = MirFnSig::new(vec![MirType::i32(), MirType::i32()], MirType::i32());
        let mut builder = MirBuilder::new("mul_test", sig);

        let a = builder.param_local(0);
        let b = builder.param_local(1);
        let result = builder.create_local(MirType::i32());

        builder.binary_op(result, BinOp::Mul, values::local(a), values::local(b));
        builder.ret(Some(values::local(result)));

        module_builder.add_function(builder.build());
        let module = module_builder.build();

        let mut backend = Arm64Backend::new();
        let output = backend.generate(&module).unwrap();

        let code = output.as_string().unwrap();
        assert!(code.contains("mul x0, x0, x1"));
    }

    #[test]
    fn test_arm64_backend_div() {
        let mut module_builder = MirModuleBuilder::new("test");

        let sig = MirFnSig::new(vec![MirType::i32(), MirType::i32()], MirType::i32());
        let mut builder = MirBuilder::new("div_test", sig);

        let a = builder.param_local(0);
        let b = builder.param_local(1);
        let result = builder.create_local(MirType::i32());

        builder.binary_op(result, BinOp::Div, values::local(a), values::local(b));
        builder.ret(Some(values::local(result)));

        module_builder.add_function(builder.build());
        let module = module_builder.build();

        let mut backend = Arm64Backend::new();
        let output = backend.generate(&module).unwrap();

        let code = output.as_string().unwrap();
        assert!(code.contains("sdiv x0, x0, x1"));
    }

    #[test]
    fn test_arm64_backend_comparison() {
        let mut module_builder = MirModuleBuilder::new("test");

        let sig = MirFnSig::new(vec![MirType::i32(), MirType::i32()], MirType::Bool);
        let mut builder = MirBuilder::new("eq_test", sig);

        let a = builder.param_local(0);
        let b = builder.param_local(1);
        let result = builder.create_local(MirType::Bool);

        builder.binary_op(result, BinOp::Eq, values::local(a), values::local(b));
        builder.ret(Some(values::local(result)));

        module_builder.add_function(builder.build());
        let module = module_builder.build();

        let mut backend = Arm64Backend::new();
        let output = backend.generate(&module).unwrap();

        let code = output.as_string().unwrap();
        assert!(code.contains("cmp x0, x1"));
        assert!(code.contains("cset x0, eq"));
    }

    #[test]
    fn test_arm64_backend_unary_neg() {
        let mut module_builder = MirModuleBuilder::new("test");

        let sig = MirFnSig::new(vec![MirType::i32()], MirType::i32());
        let mut builder = MirBuilder::new("negate", sig);

        let a = builder.param_local(0);
        let result = builder.create_local(MirType::i32());

        builder.unary_op(result, UnaryOp::Neg, values::local(a));
        builder.ret(Some(values::local(result)));

        module_builder.add_function(builder.build());
        let module = module_builder.build();

        let mut backend = Arm64Backend::new();
        let output = backend.generate(&module).unwrap();

        let code = output.as_string().unwrap();
        assert!(code.contains("neg x0, x0"));
    }

    #[test]
    fn test_arm64_backend_target() {
        let backend = Arm64Backend::new();
        assert_eq!(backend.target(), Target::Arm64);
    }

    #[test]
    fn test_arm64_backend_type_size() {
        let backend = Arm64Backend::new();

        assert_eq!(backend.type_size(&MirType::Void), 0);
        assert_eq!(backend.type_size(&MirType::Bool), 1);
        assert_eq!(backend.type_size(&MirType::i8()), 1);
        assert_eq!(backend.type_size(&MirType::i16()), 2);
        assert_eq!(backend.type_size(&MirType::i32()), 4);
        assert_eq!(backend.type_size(&MirType::i64()), 8);
        assert_eq!(backend.type_size(&MirType::f32()), 4);
        assert_eq!(backend.type_size(&MirType::f64()), 8);
        assert_eq!(
            backend.type_size(&MirType::Ptr(Box::new(MirType::i32()))),
            8
        );
    }

    // =========================================================================
    // Machine Code Generation Tests
    // =========================================================================

    #[test]
    fn test_arm64_machine_code_simple() {
        let mut module_builder = MirModuleBuilder::new("test");

        let sig = MirFnSig::new(vec![], MirType::i32());
        let mut builder = MirBuilder::new("ret42", sig);
        builder.ret(Some(values::i32(42)));

        module_builder.add_function(builder.build());
        let module = module_builder.build();

        let mut backend = Arm64Backend::new().with_machine_code();
        let output = backend.generate(&module).unwrap();

        // Check that we got object output
        assert_eq!(output.format, OutputFormat::Object);
        assert!(!output.data.is_empty());

        // ARM64 instructions are 4 bytes each
        assert!(output.data.len() % 4 == 0);
    }

    #[test]
    fn test_arm64_machine_code_add() {
        let mut module_builder = MirModuleBuilder::new("test");

        let sig = MirFnSig::new(vec![MirType::i64(), MirType::i64()], MirType::i64());
        let mut builder = MirBuilder::new("add", sig);

        let a = builder.param_local(0);
        let b = builder.param_local(1);
        let result = builder.create_local(MirType::i64());

        builder.binary_op(result, BinOp::Add, values::local(a), values::local(b));
        builder.ret(Some(values::local(result)));

        module_builder.add_function(builder.build());
        let module = module_builder.build();

        let mut backend = Arm64Backend::new().with_machine_code();
        let output = backend.generate(&module).unwrap();

        assert_eq!(output.format, OutputFormat::Object);
        assert!(!output.data.is_empty());
    }

    #[test]
    fn test_arm64_machine_code_with_branch() {
        let mut module_builder = MirModuleBuilder::new("test");

        let sig = MirFnSig::new(vec![MirType::Bool], MirType::i32());
        let mut builder = MirBuilder::new("branch", sig);

        let cond = builder.param_local(0);
        let then_block = builder.create_block();
        let else_block = builder.create_block();

        builder.branch(values::local(cond), then_block, else_block);

        builder.switch_to_block(then_block);
        builder.ret(Some(values::i32(1)));

        builder.switch_to_block(else_block);
        builder.ret(Some(values::i32(0)));

        module_builder.add_function(builder.build());
        let module = module_builder.build();

        let mut backend = Arm64Backend::new().with_machine_code();
        let output = backend.generate(&module).unwrap();

        assert_eq!(output.format, OutputFormat::Object);
        assert!(!output.data.is_empty());
    }

    #[test]
    fn test_arm64_machine_code_prologue_epilogue() {
        let mut module_builder = MirModuleBuilder::new("test");

        let sig = MirFnSig::new(vec![], MirType::Void);
        let mut builder = MirBuilder::new("noop", sig);
        builder.ret_void();

        module_builder.add_function(builder.build());
        let module = module_builder.build();

        let mut backend = Arm64Backend::new().with_machine_code();
        let output = backend.generate(&module).unwrap();

        // Should have prologue + epilogue + ret
        // Minimum: stp + mov + ldp + ret = 4 instructions = 16 bytes
        assert!(output.data.len() >= 16);

        // Last instruction should be RET (0xD65F03C0)
        let last_insn =
            u32::from_le_bytes(output.data[output.data.len() - 4..].try_into().unwrap());
        assert_eq!(last_insn, 0xD65F03C0);
    }

    #[test]
    fn test_arm64_machine_code_binary_ops() {
        let ops = [
            (BinOp::Add, "add"),
            (BinOp::Sub, "sub"),
            (BinOp::Mul, "mul"),
            (BinOp::BitAnd, "and"),
            (BinOp::BitOr, "or"),
            (BinOp::BitXor, "xor"),
        ];

        for (op, name) in ops {
            let mut module_builder = MirModuleBuilder::new("test");

            let sig = MirFnSig::new(vec![MirType::i64(), MirType::i64()], MirType::i64());
            let mut builder = MirBuilder::new(name, sig);

            let a = builder.param_local(0);
            let b = builder.param_local(1);
            let result = builder.create_local(MirType::i64());

            builder.binary_op(result, op, values::local(a), values::local(b));
            builder.ret(Some(values::local(result)));

            module_builder.add_function(builder.build());
            let module = module_builder.build();

            let mut backend = Arm64Backend::new().with_machine_code();
            let output = backend.generate(&module).unwrap();

            assert_eq!(output.format, OutputFormat::Object);
            assert!(!output.data.is_empty(), "Failed for op: {}", name);
        }
    }

    #[test]
    fn test_arm64_machine_code_comparison() {
        let mut module_builder = MirModuleBuilder::new("test");

        let sig = MirFnSig::new(vec![MirType::i64(), MirType::i64()], MirType::Bool);
        let mut builder = MirBuilder::new("eq", sig);

        let a = builder.param_local(0);
        let b = builder.param_local(1);
        let result = builder.create_local(MirType::Bool);

        builder.binary_op(result, BinOp::Eq, values::local(a), values::local(b));
        builder.ret(Some(values::local(result)));

        module_builder.add_function(builder.build());
        let module = module_builder.build();

        let mut backend = Arm64Backend::new().with_machine_code();
        let output = backend.generate(&module).unwrap();

        assert_eq!(output.format, OutputFormat::Object);
        assert!(!output.data.is_empty());
    }

    #[test]
    fn test_arm64_output_mode() {
        let backend = Arm64Backend::new();
        assert_eq!(backend.mode, Arm64OutputMode::Assembly);

        let backend = Arm64Backend::new().with_machine_code();
        assert_eq!(backend.mode, Arm64OutputMode::MachineCode);
        assert!(backend.encoder.is_some());
    }

    #[test]
    fn test_arm64_machine_code_unary_neg() {
        let mut module_builder = MirModuleBuilder::new("test");

        let sig = MirFnSig::new(vec![MirType::i64()], MirType::i64());
        let mut builder = MirBuilder::new("negate", sig);

        let a = builder.param_local(0);
        let result = builder.create_local(MirType::i64());

        builder.unary_op(result, UnaryOp::Neg, values::local(a));
        builder.ret(Some(values::local(result)));

        module_builder.add_function(builder.build());
        let module = module_builder.build();

        let mut backend = Arm64Backend::new().with_machine_code();
        let output = backend.generate(&module).unwrap();

        assert_eq!(output.format, OutputFormat::Object);
        assert!(!output.data.is_empty());
    }

    #[test]
    fn test_arm64_machine_code_multiple_functions() {
        let mut module_builder = MirModuleBuilder::new("test");

        // First function
        let sig1 = MirFnSig::new(vec![], MirType::i32());
        let mut builder1 = MirBuilder::new("func1", sig1);
        builder1.ret(Some(values::i32(1)));
        module_builder.add_function(builder1.build());

        // Second function
        let sig2 = MirFnSig::new(vec![], MirType::i32());
        let mut builder2 = MirBuilder::new("func2", sig2);
        builder2.ret(Some(values::i32(2)));
        module_builder.add_function(builder2.build());

        let module = module_builder.build();

        let mut backend = Arm64Backend::new().with_machine_code();
        let output = backend.generate(&module).unwrap();

        assert_eq!(output.format, OutputFormat::Object);
        // Should have code for both functions
        assert!(output.data.len() >= 32); // At least 8 instructions
    }

    #[test]
    fn test_arm64_machine_code_div() {
        let mut module_builder = MirModuleBuilder::new("test");

        let sig = MirFnSig::new(vec![MirType::i64(), MirType::i64()], MirType::i64());
        let mut builder = MirBuilder::new("div", sig);

        let a = builder.param_local(0);
        let b = builder.param_local(1);
        let result = builder.create_local(MirType::i64());

        builder.binary_op(result, BinOp::Div, values::local(a), values::local(b));
        builder.ret(Some(values::local(result)));

        module_builder.add_function(builder.build());
        let module = module_builder.build();

        let mut backend = Arm64Backend::new().with_machine_code();
        let output = backend.generate(&module).unwrap();

        assert_eq!(output.format, OutputFormat::Object);
        assert!(!output.data.is_empty());
    }

    #[test]
    fn test_arm64_machine_code_rem() {
        let mut module_builder = MirModuleBuilder::new("test");

        let sig = MirFnSig::new(vec![MirType::i64(), MirType::i64()], MirType::i64());
        let mut builder = MirBuilder::new("rem", sig);

        let a = builder.param_local(0);
        let b = builder.param_local(1);
        let result = builder.create_local(MirType::i64());

        builder.binary_op(result, BinOp::Rem, values::local(a), values::local(b));
        builder.ret(Some(values::local(result)));

        module_builder.add_function(builder.build());
        let module = module_builder.build();

        let mut backend = Arm64Backend::new().with_machine_code();
        let output = backend.generate(&module).unwrap();

        assert_eq!(output.format, OutputFormat::Object);
        // Remainder uses SDIV + MSUB, so should have more code
        assert!(output.data.len() >= 20);
    }
}
