//! AOT Compiler - Ahead-of-time compilation to native object files.
//!
//! Uses cranelift-object to compile JVM bytecode to .o files that can be
//! linked into native executables. Supports LTO-style optimizations via
//! Cranelift's egraph-based optimizer and configurable optimization levels.

use crate::class_file::{ClassFile, MethodInfo};
use crate::jit::JitError;
use cranelift::codegen::ir::{AbiParam, Signature};
use cranelift::codegen::isa::CallConv;
use cranelift::prelude::*;
use cranelift_frontend::{FunctionBuilder, FunctionBuilderContext};
use cranelift_module::{Linkage, Module};
use cranelift_object::{ObjectBuilder, ObjectModule};
use std::path::Path;

/// Configuration for AOT compilation, including LTO-style optimizations.
#[derive(Debug, Clone)]
pub struct AotCompilationConfig {
    /// Optimization level: "none", "speed", "speed_and_size"
    pub opt_level: String,
    /// Enable alias analysis for better load/store optimization
    pub enable_alias_analysis: bool,
}

impl Default for AotCompilationConfig {
    fn default() -> Self {
        Self {
            opt_level: "speed".to_string(),
            enable_alias_analysis: true,
        }
    }
}

/// Build Cranelift IR for a method and define it in the module.
fn build_method_ir(
    method: &MethodInfo,
    module: &mut dyn Module,
    func_name: &str,
    builder_ctx: &mut FunctionBuilderContext,
) -> Result<(), JitError> {
    let code_attr = method
        .attributes
        .iter()
        .find(|a| a.info.len() >= 8)
        .ok_or_else(|| JitError::InvalidMethod("No Code attribute".to_string()))?;

    let code_len = ((code_attr.info[4] as usize) << 24)
        | ((code_attr.info[5] as usize) << 16)
        | ((code_attr.info[6] as usize) << 8)
        | (code_attr.info[7] as usize);
    let bytecode = code_attr.info.get(8..8 + code_len).unwrap_or(&[]);

    let ptr_type = module.target_config().pointer_type();
    let mut sig = Signature::new(CallConv::SystemV);
    sig.params.push(AbiParam::new(ptr_type));
    sig.params.push(AbiParam::new(ptr_type));
    sig.returns.push(AbiParam::new(types::I32));

    macro_rules! decl_helper {
        ($name:expr, ret $ret:expr, params [$($param:expr),*]) => {{
            let mut s = Signature::new(CallConv::SystemV);
            $( s.params.push(AbiParam::new($param)); )*
            s.returns.push(AbiParam::new($ret));
            module.declare_function($name, Linkage::Import, &s)
                .map_err(|e| JitError::CompilationFailed(e.to_string()))?
        }};
        ($name:expr, void, params [$($param:expr),*]) => {{
            let mut s = Signature::new(CallConv::SystemV);
            $( s.params.push(AbiParam::new($param)); )*
            module.declare_function($name, Linkage::Import, &s)
                .map_err(|e| JitError::CompilationFailed(e.to_string()))?
        }};
    }

    let get_int_fn    = decl_helper!("jvmrs_frame_get_local_int",    ret types::I32,  params [ptr_type, types::I32]);
    let push_int_fn   = decl_helper!("jvmrs_frame_push_int",       void,            params [ptr_type, types::I32]);
    let get_long_fn   = decl_helper!("jvmrs_frame_get_local_long", ret types::I64,  params [ptr_type, types::I32]);
    let push_long_fn  = decl_helper!("jvmrs_frame_push_long",      void,            params [ptr_type, types::I64]);
    let get_float_fn  = decl_helper!("jvmrs_frame_get_local_float", ret types::F32,  params [ptr_type, types::I32]);
    let push_float_fn = decl_helper!("jvmrs_frame_push_float",     void,            params [ptr_type, types::F32]);
    let set_float_fn  = decl_helper!("jvmrs_frame_set_local_float", void,           params [ptr_type, types::I32, types::F32]);
    let get_double_fn = decl_helper!("jvmrs_frame_get_local_double", ret types::F64, params [ptr_type, types::I32]);
    let push_double_fn= decl_helper!("jvmrs_frame_push_double",    void,            params [ptr_type, types::F64]);
    let set_double_fn = decl_helper!("jvmrs_frame_set_local_double", void,           params [ptr_type, types::I32, types::F64]);
    let get_ref_fn    = decl_helper!("jvmrs_frame_get_local_ref",  ret types::I32,  params [ptr_type, types::I32]);
    let push_ref_fn   = decl_helper!("jvmrs_frame_push_ref",       void,            params [ptr_type, types::I32]);
    let set_ref_fn    = decl_helper!("jvmrs_frame_set_local_ref",  void,            params [ptr_type, types::I32, types::I32]);
    let set_int_fn    = decl_helper!("jvmrs_frame_set_local_int",  void,            params [ptr_type, types::I32, types::I32]);
    let set_long_fn   = decl_helper!("jvmrs_frame_set_local_long", void,            params [ptr_type, types::I32, types::I64]);
    let inc_int_fn    = decl_helper!("jvmrs_frame_inc_local_int",  void,            params [ptr_type, types::I32, types::I32]);

    let func_id = module
        .declare_function(func_name, Linkage::Export, &sig)
        .map_err(|e| JitError::CompilationFailed(e.to_string()))?;

    let mut ctx = module.make_context();
    ctx.func.signature = sig.clone();

    let get_int_ref    = module.declare_func_in_func(get_int_fn,    &mut ctx.func);
    let push_int_ref   = module.declare_func_in_func(push_int_fn,   &mut ctx.func);
    let set_int_ref    = module.declare_func_in_func(set_int_fn,    &mut ctx.func);
    let get_long_ref   = module.declare_func_in_func(get_long_fn,   &mut ctx.func);
    let push_long_ref  = module.declare_func_in_func(push_long_fn,  &mut ctx.func);
    let set_long_ref   = module.declare_func_in_func(set_long_fn,  &mut ctx.func);
    let get_float_ref  = module.declare_func_in_func(get_float_fn,  &mut ctx.func);
    let push_float_ref = module.declare_func_in_func(push_float_fn, &mut ctx.func);
    let set_float_ref  = module.declare_func_in_func(set_float_fn,  &mut ctx.func);
    let get_double_ref = module.declare_func_in_func(get_double_fn, &mut ctx.func);
    let push_double_ref= module.declare_func_in_func(push_double_fn,&mut ctx.func);
    let set_double_ref = module.declare_func_in_func(set_double_fn, &mut ctx.func);
    let get_ref_ref    = module.declare_func_in_func(get_ref_fn,    &mut ctx.func);
    let push_ref_ref   = module.declare_func_in_func(push_ref_fn,   &mut ctx.func);
    let set_ref_ref    = module.declare_func_in_func(set_ref_fn,    &mut ctx.func);
    let inc_int_ref    = module.declare_func_in_func(inc_int_fn,    &mut ctx.func);

    let mut builder = FunctionBuilder::new(&mut ctx.func, builder_ctx);

    let entry = builder.create_block();
    builder.append_block_params_for_function_params(entry);
    builder.switch_to_block(entry);
    builder.seal_block(entry);

    let params = builder.block_params(entry);
    let frame_param = params[1];

    let i32_type = types::I32;
    let i64_type = types::I64;
    let f32_type = types::F32;
    let f64_type = types::F64;
    let mut stack: Vec<cranelift::prelude::Value> = Vec::new();

    let mut pc = 0usize;
    let mut compiled = false;

    while pc < bytecode.len() {
        let opcode = bytecode[pc];
        pc += 1;

        match opcode {
            // ----- Constants -----
            0x00 => {} // nop
            0x01 => stack.push(builder.ins().iconst(i32_type, 0)), // aconst_null
            0x02 => stack.push(builder.ins().iconst(i32_type, -1)),
            0x03 => stack.push(builder.ins().iconst(i32_type, 0)),
            0x04 => stack.push(builder.ins().iconst(i32_type, 1)),
            0x05 => stack.push(builder.ins().iconst(i32_type, 2)),
            0x06 => stack.push(builder.ins().iconst(i32_type, 3)),
            0x07 => stack.push(builder.ins().iconst(i32_type, 4)),
            0x08 => stack.push(builder.ins().iconst(i32_type, 5)),
            0x09 => stack.push(builder.ins().iconst(i64_type, 0)),
            0x0a => stack.push(builder.ins().iconst(i64_type, 1)),
            0x0b => stack.push(builder.ins().f32const(0.0)),
            0x0c => stack.push(builder.ins().f32const(1.0)),
            0x0d => stack.push(builder.ins().f32const(2.0)),
            0x0e => stack.push(builder.ins().f64const(0.0)),
            0x0f => stack.push(builder.ins().f64const(1.0)),
            0x10 => { // bipush
                if pc < bytecode.len() {
                    let v = bytecode[pc] as i8 as i32;
                    pc += 1;
                    stack.push(builder.ins().iconst(i32_type, v as i64));
                }
            }
            0x11 => { // sipush
                if pc + 1 < bytecode.len() {
                    let v = ((bytecode[pc] as i16) << 8) | (bytecode[pc + 1] as i16);
                    pc += 2;
                    stack.push(builder.ins().iconst(i32_type, v as i64));
                }
            }
            0x12 | 0x13 => { pc += 1; stack.push(builder.ins().iconst(i32_type, 0)); }

            // ----- Loads -----
            0x1a..=0x1d => { // iload_0..3
                let idx = builder.ins().iconst(i32_type, (opcode - 0x1a) as i64);
                let call = builder.ins().call(get_int_ref, &[frame_param, idx]);
                stack.push(builder.inst_results(call)[0]);
            }
            0x1e..=0x21 => { // lload_0..3
                let idx = builder.ins().iconst(i32_type, (opcode - 0x1e) as i64);
                let call = builder.ins().call(get_long_ref, &[frame_param, idx]);
                stack.push(builder.inst_results(call)[0]);
            }
            0x22..=0x25 => { // fload_0..3
                let idx = builder.ins().iconst(i32_type, (opcode - 0x22) as i64);
                let call = builder.ins().call(get_float_ref, &[frame_param, idx]);
                stack.push(builder.inst_results(call)[0]);
            }
            0x26..=0x29 => { // dload_0..3
                let idx = builder.ins().iconst(i32_type, (opcode - 0x26) as i64);
                let call = builder.ins().call(get_double_ref, &[frame_param, idx]);
                stack.push(builder.inst_results(call)[0]);
            }
            0x2a..=0x2d => { // aload_0..3
                let idx = builder.ins().iconst(i32_type, (opcode - 0x2a) as i64);
                let call = builder.ins().call(get_ref_ref, &[frame_param, idx]);
                stack.push(builder.inst_results(call)[0]);
            }
            0x15 => { // iload
                if pc < bytecode.len() {
                    let idx = builder.ins().iconst(i32_type, bytecode[pc] as i64);
                    pc += 1;
                    let call = builder.ins().call(get_int_ref, &[frame_param, idx]);
                    stack.push(builder.inst_results(call)[0]);
                }
            }
            0x19 => { // aload
                if pc < bytecode.len() {
                    let idx = builder.ins().iconst(i32_type, bytecode[pc] as i64);
                    pc += 1;
                    let call = builder.ins().call(get_ref_ref, &[frame_param, idx]);
                    stack.push(builder.inst_results(call)[0]);
                }
            }

            // ----- Stores -----
            0x3b..=0x3e => { // istore_0..3
                let idx = builder.ins().iconst(i32_type, (opcode - 0x3b) as i64);
                if let Some(v) = stack.pop() {
                    builder.ins().call(set_int_ref, &[frame_param, idx, v]);
                }
            }
            0x3f..=0x42 => { // lstore_0..3
                let idx = builder.ins().iconst(i32_type, (opcode - 0x3f) as i64);
                if let Some(v) = stack.pop() {
                    builder.ins().call(set_long_ref, &[frame_param, idx, v]);
                }
            }
            0x43..=0x46 => { // fstore_0..3
                let idx = builder.ins().iconst(i32_type, (opcode - 0x43) as i64);
                if let Some(v) = stack.pop() {
                    builder.ins().call(set_float_ref, &[frame_param, idx, v]);
                }
            }
            0x47..=0x4a => { // dstore_0..3
                let idx = builder.ins().iconst(i32_type, (opcode - 0x47) as i64);
                if let Some(v) = stack.pop() {
                    builder.ins().call(set_double_ref, &[frame_param, idx, v]);
                }
            }
            0x4b..=0x4e => { // astore_0..3
                let idx = builder.ins().iconst(i32_type, (opcode - 0x4b) as i64);
                if let Some(v) = stack.pop() {
                    builder.ins().call(set_ref_ref, &[frame_param, idx, v]);
                }
            }
            0x36 => { // istore
                if pc < bytecode.len() {
                    let idx = builder.ins().iconst(i32_type, bytecode[pc] as i64);
                    pc += 1;
                    if let Some(v) = stack.pop() {
                        builder.ins().call(set_int_ref, &[frame_param, idx, v]);
                    }
                }
            }
            0x37 => { // lstore
                if pc < bytecode.len() {
                    let idx = builder.ins().iconst(i32_type, bytecode[pc] as i64);
                    pc += 1;
                    if let Some(v) = stack.pop() {
                        builder.ins().call(set_long_ref, &[frame_param, idx, v]);
                    }
                }
            }
            0x38 => { // fstore
                if pc < bytecode.len() {
                    let idx = builder.ins().iconst(i32_type, bytecode[pc] as i64);
                    pc += 1;
                    if let Some(v) = stack.pop() {
                        builder.ins().call(set_float_ref, &[frame_param, idx, v]);
                    }
                }
            }
            0x39 => { // dstore
                if pc < bytecode.len() {
                    let idx = builder.ins().iconst(i32_type, bytecode[pc] as i64);
                    pc += 1;
                    if let Some(v) = stack.pop() {
                        builder.ins().call(set_double_ref, &[frame_param, idx, v]);
                    }
                }
            }
            0x3a => { // astore
                if pc < bytecode.len() {
                    let idx = builder.ins().iconst(i32_type, bytecode[pc] as i64);
                    pc += 1;
                    if let Some(v) = stack.pop() {
                        builder.ins().call(set_ref_ref, &[frame_param, idx, v]);
                    }
                }
            }

            // ----- Stack -----
            0x57 => { stack.pop(); }
            0x58 => { stack.pop(); stack.pop(); }
            0x59 => {
                let v = stack.last().copied().unwrap_or_else(|| builder.ins().iconst(i32_type, 0));
                stack.push(v);
            }
            0x5f => {
                let v1 = stack.pop().unwrap_or_else(|| builder.ins().iconst(i32_type, 0));
                let v2 = stack.pop().unwrap_or_else(|| builder.ins().iconst(i32_type, 0));
                stack.push(v1);
                stack.push(v2);
            }

            // ----- Integer arithmetic -----
            0x60 => {
                let b = stack.pop().unwrap_or_else(|| builder.ins().iconst(i32_type, 0));
                let a = stack.pop().unwrap_or_else(|| builder.ins().iconst(i32_type, 0));
                stack.push(builder.ins().iadd(a, b));
            }
            0x64 => {
                let b = stack.pop().unwrap_or_else(|| builder.ins().iconst(i32_type, 0));
                let a = stack.pop().unwrap_or_else(|| builder.ins().iconst(i32_type, 0));
                stack.push(builder.ins().isub(a, b));
            }
            0x68 => {
                let b = stack.pop().unwrap_or_else(|| builder.ins().iconst(i32_type, 0));
                let a = stack.pop().unwrap_or_else(|| builder.ins().iconst(i32_type, 0));
                stack.push(builder.ins().imul(a, b));
            }
            0x6c => {
                let b = stack.pop().unwrap_or_else(|| builder.ins().iconst(i32_type, 0));
                let a = stack.pop().unwrap_or_else(|| builder.ins().iconst(i32_type, 0));
                stack.push(builder.ins().sdiv(a, b));
            }
            0x70 => {
                let b = stack.pop().unwrap_or_else(|| builder.ins().iconst(i32_type, 0));
                let a = stack.pop().unwrap_or_else(|| builder.ins().iconst(i32_type, 0));
                stack.push(builder.ins().srem(a, b));
            }
            0x74 => {
                let a = stack.pop().unwrap_or_else(|| builder.ins().iconst(i32_type, 0));
                stack.push(builder.ins().ineg(a));
            }
            0x61 => {
                let b = stack.pop().unwrap_or_else(|| builder.ins().iconst(i64_type, 0));
                let a = stack.pop().unwrap_or_else(|| builder.ins().iconst(i64_type, 0));
                stack.push(builder.ins().iadd(a, b));
            }
            0x65 => {
                let b = stack.pop().unwrap_or_else(|| builder.ins().iconst(i64_type, 0));
                let a = stack.pop().unwrap_or_else(|| builder.ins().iconst(i64_type, 0));
                stack.push(builder.ins().isub(a, b));
            }
            0x69 => {
                let b = stack.pop().unwrap_or_else(|| builder.ins().iconst(i64_type, 0));
                let a = stack.pop().unwrap_or_else(|| builder.ins().iconst(i64_type, 0));
                stack.push(builder.ins().imul(a, b));
            }
            0x6d => {
                let b = stack.pop().unwrap_or_else(|| builder.ins().iconst(i64_type, 0));
                let a = stack.pop().unwrap_or_else(|| builder.ins().iconst(i64_type, 0));
                stack.push(builder.ins().sdiv(a, b));
            }
            0x71 => {
                let b = stack.pop().unwrap_or_else(|| builder.ins().iconst(i64_type, 0));
                let a = stack.pop().unwrap_or_else(|| builder.ins().iconst(i64_type, 0));
                stack.push(builder.ins().srem(a, b));
            }
            0x75 => {
                let a = stack.pop().unwrap_or_else(|| builder.ins().iconst(i64_type, 0));
                stack.push(builder.ins().ineg(a));
            }

            // ----- Float arithmetic -----
            0x62 => {
                let b = stack.pop().unwrap_or_else(|| builder.ins().f32const(0.0));
                let a = stack.pop().unwrap_or_else(|| builder.ins().f32const(0.0));
                stack.push(builder.ins().fadd(a, b));
            }
            0x66 => {
                let b = stack.pop().unwrap_or_else(|| builder.ins().f32const(0.0));
                let a = stack.pop().unwrap_or_else(|| builder.ins().f32const(0.0));
                stack.push(builder.ins().fsub(a, b));
            }
            0x6a => {
                let b = stack.pop().unwrap_or_else(|| builder.ins().f32const(0.0));
                let a = stack.pop().unwrap_or_else(|| builder.ins().f32const(0.0));
                stack.push(builder.ins().fmul(a, b));
            }
            0x6e => {
                let b = stack.pop().unwrap_or_else(|| builder.ins().f32const(0.0));
                let a = stack.pop().unwrap_or_else(|| builder.ins().f32const(0.0));
                stack.push(builder.ins().fdiv(a, b));
            }

            // ----- Double arithmetic -----
            0x63 => {
                let b = stack.pop().unwrap_or_else(|| builder.ins().f64const(0.0));
                let a = stack.pop().unwrap_or_else(|| builder.ins().f64const(0.0));
                stack.push(builder.ins().fadd(a, b));
            }
            0x67 => {
                let b = stack.pop().unwrap_or_else(|| builder.ins().f64const(0.0));
                let a = stack.pop().unwrap_or_else(|| builder.ins().f64const(0.0));
                stack.push(builder.ins().fsub(a, b));
            }
            0x6b => {
                let b = stack.pop().unwrap_or_else(|| builder.ins().f64const(0.0));
                let a = stack.pop().unwrap_or_else(|| builder.ins().f64const(0.0));
                stack.push(builder.ins().fmul(a, b));
            }
            0x6f => {
                let b = stack.pop().unwrap_or_else(|| builder.ins().f64const(0.0));
                let a = stack.pop().unwrap_or_else(|| builder.ins().f64const(0.0));
                stack.push(builder.ins().fdiv(a, b));
            }

            // ----- Conversions -----
            0x85 => {
                let a = stack.pop().unwrap_or_else(|| builder.ins().iconst(i32_type, 0));
                stack.push(builder.ins().sextend(i64_type, a));
            }
            0x86 => {
                let a = stack.pop().unwrap_or_else(|| builder.ins().iconst(i32_type, 0));
                stack.push(builder.ins().fcvt_from_sint(f32_type, a));
            }
            0x87 => {
                let a = stack.pop().unwrap_or_else(|| builder.ins().iconst(i32_type, 0));
                stack.push(builder.ins().fcvt_from_sint(f64_type, a));
            }
            0x88 => {
                let a = stack.pop().unwrap_or_else(|| builder.ins().iconst(i64_type, 0));
                stack.push(builder.ins().ireduce(i32_type, a));
            }
            0x8b => {
                let a = stack.pop().unwrap_or_else(|| builder.ins().f32const(0.0));
                stack.push(builder.ins().fcvt_to_sint_sat(i32_type, a));
            }
            0x8d => {
                let a = stack.pop().unwrap_or_else(|| builder.ins().f32const(0.0));
                stack.push(builder.ins().fpromote(f64_type, a));
            }
            0x8e => {
                let a = stack.pop().unwrap_or_else(|| builder.ins().f64const(0.0));
                stack.push(builder.ins().fcvt_to_sint_sat(i32_type, a));
            }
            0x90 => {
                let a = stack.pop().unwrap_or_else(|| builder.ins().f64const(0.0));
                stack.push(builder.ins().fdemote(f32_type, a));
            }

            // ----- Increment -----
            0x84 => {
                if pc + 1 < bytecode.len() {
                    let idx = builder.ins().iconst(i32_type, bytecode[pc] as i64);
                    let delta = builder.ins().iconst(i32_type, bytecode[pc + 1] as i8 as i64);
                    pc += 2;
                    builder.ins().call(inc_int_ref, &[frame_param, idx, delta]);
                }
            }

            // ----- Returns -----
            0xb1 => {
                let ret_val = builder.ins().iconst(i32_type, 0);
                builder.ins().return_(&[ret_val]);
                compiled = true;
                break;
            }
            0xac => {
                let result = stack.pop().unwrap_or_else(|| builder.ins().iconst(i32_type, 0));
                builder.ins().call(push_int_ref, &[frame_param, result]);
                let ret_val = builder.ins().iconst(i32_type, 0);
                builder.ins().return_(&[ret_val]);
                compiled = true;
                break;
            }
            0xad => {
                let result = stack.pop().unwrap_or_else(|| builder.ins().iconst(i64_type, 0));
                builder.ins().call(push_long_ref, &[frame_param, result]);
                let ret_val = builder.ins().iconst(i32_type, 0);
                builder.ins().return_(&[ret_val]);
                compiled = true;
                break;
            }
            0xae => {
                let result = stack.pop().unwrap_or_else(|| builder.ins().f32const(0.0));
                builder.ins().call(push_float_ref, &[frame_param, result]);
                let ret_val = builder.ins().iconst(i32_type, 0);
                builder.ins().return_(&[ret_val]);
                compiled = true;
                break;
            }
            0xaf => {
                let result = stack.pop().unwrap_or_else(|| builder.ins().f64const(0.0));
                builder.ins().call(push_double_ref, &[frame_param, result]);
                let ret_val = builder.ins().iconst(i32_type, 0);
                builder.ins().return_(&[ret_val]);
                compiled = true;
                break;
            }
            0xb0 => {
                let result = stack.pop().unwrap_or_else(|| builder.ins().iconst(i32_type, 0));
                builder.ins().call(push_ref_ref, &[frame_param, result]);
                let ret_val = builder.ins().iconst(i32_type, 0);
                builder.ins().return_(&[ret_val]);
                compiled = true;
                break;
            }

            _ => {
                return Err(JitError::UnsupportedInstruction(format!(
                    "opcode 0x{:02x} at pc {}",
                    opcode, pc
                )));
            }
        }
    }

    if !compiled {
        let ret_val = builder.ins().iconst(i32_type, 0);
        builder.ins().return_(&[ret_val]);
    }

    builder.finalize();

    module
        .define_function(func_id, &mut ctx)
        .map_err(|e| JitError::CompilationFailed(e.to_string()))?;

    Ok(())
}

/// Compile a class to a native object file (.o) using default optimizations.
pub fn compile_class_to_object(class: &ClassFile, output_path: &Path) -> Result<(), JitError> {
    compile_class_to_object_with_config(class, output_path, &AotCompilationConfig::default())
}

/// Compile a class to a native object file (.o) with configurable optimizations.
pub fn compile_class_to_object_with_config(
    class: &ClassFile,
    output_path: &Path,
    config: &AotCompilationConfig,
) -> Result<(), JitError> {
    let class_name = class
        .get_class_name()
        .unwrap_or_else(|| "Unknown".to_string());
    let safe_name = class_name.replace("/", "_").replace(".", "_");

    let mut flag_builder = cranelift::codegen::settings::builder();
    flag_builder
        .set("use_colocated_libcalls", "false")
        .map_err(|e| JitError::CompilationFailed(e.to_string()))?;
    flag_builder
        .set("is_pic", "true")
        .map_err(|e| JitError::CompilationFailed(e.to_string()))?;
    // LTO-style optimizations
    flag_builder
        .set("opt_level", &config.opt_level)
        .map_err(|e| JitError::CompilationFailed(e.to_string()))?;
    flag_builder
        .set("enable_alias_analysis", if config.enable_alias_analysis { "true" } else { "false" })
        .map_err(|e| JitError::CompilationFailed(e.to_string()))?;

    let isa = cranelift_native::builder()
        .map_err(|e| JitError::CompilationFailed(format!("ISA: {}", e)))?
        .finish(cranelift::codegen::settings::Flags::new(flag_builder))
        .map_err(|e| JitError::CompilationFailed(e.to_string()))?;

    let builder = ObjectBuilder::new(
        isa,
        safe_name.as_str(),
        cranelift_module::default_libcall_names(),
    )
    .map_err(|e| JitError::CompilationFailed(format!("ObjectBuilder: {}", e)))?;
    let mut module = ObjectModule::new(builder);
    let mut builder_ctx = FunctionBuilderContext::new();

    let mut compiled = 0usize;

    for method in &class.methods {
        let method_name = class
            .get_string(method.name_index)
            .unwrap_or_else(|| "unknown".to_string());

        if method.access_flags & 0x0400 != 0 || method.access_flags & 0x0100 != 0 {
            continue;
        }

        if method.attributes.iter().all(|a| a.info.len() < 8) {
            continue;
        }

        let func_name = format!(
            "{}_{}",
            safe_name,
            method_name.replace("<", "_").replace(">", "_")
        );

        if build_method_ir(method, &mut module, &func_name, &mut builder_ctx).is_ok() {
            compiled += 1;
        }
    }

    if compiled == 0 {
        return Err(JitError::CompilationFailed(
            "No compilable methods found".to_string(),
        ));
    }

    let product = module.finish();
    let bytes = product
        .emit()
        .map_err(|e| JitError::LinkingError(format!("Object emission failed: {}", e)))?;

    std::fs::write(output_path, &bytes)
        .map_err(|e| JitError::LinkingError(format!("Failed to write object file: {}", e)))?;

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::class_file::{AttributeInfo, ClassFile, ConstantPoolEntry, FieldInfo, MethodInfo};

    fn make_code_attribute(bytecode: &[u8]) -> AttributeInfo {
        let mut info = Vec::new();
        info.extend_from_slice(&[0u8, 10]); // max_stack = 10
        info.extend_from_slice(&[0u8, 10]); // max_locals = 10
        let code_len = bytecode.len() as u32;
        info.extend_from_slice(&code_len.to_be_bytes());
        info.extend_from_slice(bytecode);
        info.extend_from_slice(&[0u8, 0]); // exception_table_length = 0
        info.extend_from_slice(&[0u8, 0]); // attributes_count = 0
        AttributeInfo {
            attribute_name_index: 4, // "Code"
            info,
        }
    }

    fn make_test_class(bytecode: &[u8]) -> ClassFile {
        ClassFile {
            magic: 0xCAFEBABE,
            minor_version: 0,
            major_version: 52,
            constant_pool: vec![
                ConstantPoolEntry::ConstantUtf8 { bytes: b"Test".to_vec() },          // 1
                ConstantPoolEntry::ConstantUtf8 { bytes: b"java/lang/Object".to_vec() }, // 2
                ConstantPoolEntry::ConstantUtf8 { bytes: b"main".to_vec() },            // 3
                ConstantPoolEntry::ConstantUtf8 { bytes: b"Code".to_vec() },            // 4
                ConstantPoolEntry::ConstantUtf8 { bytes: b"()I".to_vec() },             // 5
                ConstantPoolEntry::ConstantClass { name_index: 1 },                     // 6
                ConstantPoolEntry::ConstantClass { name_index: 2 },                       // 7
                ConstantPoolEntry::ConstantNameAndType { name_index: 3, descriptor_index: 5 }, // 8
            ],
            access_flags: 0x0001,
            this_class: 6,
            super_class: 7,
            interfaces: vec![],
            fields: vec![],
            methods: vec![MethodInfo {
                access_flags: 0x0009, // public static
                name_index: 3,
                descriptor_index: 5,
                attributes: vec![make_code_attribute(bytecode)],
            }],
            attributes: vec![],
        }
    }

    #[test]
    fn test_aot_config_default() {
        let cfg = AotCompilationConfig::default();
        assert_eq!(cfg.opt_level, "speed");
        assert!(cfg.enable_alias_analysis);
    }

    #[test]
    fn test_aot_compile_simple_method() {
        // iconst_1, iconst_2, iadd, ireturn
        let bytecode = [0x04, 0x05, 0x60, 0xac];
        let class = make_test_class(&bytecode);
        let tmp = std::env::temp_dir().join("jvmrs_aot_test_simple.o");
        let result = compile_class_to_object(&class, &tmp);
        assert!(result.is_ok(), "AOT compilation failed: {:?}", result);
        let metadata = std::fs::metadata(&tmp);
        assert!(metadata.is_ok());
        assert!(metadata.unwrap().len() > 0);
        let _ = std::fs::remove_file(&tmp);
    }

    #[test]
    fn test_aot_compile_with_expanded_opcodes() {
        // bipush 42, sipush 1000, iconst_0, iadd, isub, imul, ireturn
        let bytecode = [
            0x10, 0x2a,       // bipush 42
            0x11, 0x03, 0xe8, // sipush 1000
            0x03,             // iconst_0
            0x60,             // iadd
            0x64,             // isub
            0x68,             // imul
            0xac,             // ireturn
        ];
        let class = make_test_class(&bytecode);
        let tmp = std::env::temp_dir().join("jvmrs_aot_test_expanded.o");
        let result = compile_class_to_object(&class, &tmp);
        assert!(result.is_ok(), "AOT compilation with expanded opcodes failed: {:?}", result);
        let _ = std::fs::remove_file(&tmp);
    }

    #[test]
    fn test_aot_compile_with_no_optimizations() {
        let bytecode = [0x04, 0x05, 0x60, 0xac];
        let class = make_test_class(&bytecode);
        let tmp = std::env::temp_dir().join("jvmrs_aot_test_no_opt.o");
        let config = AotCompilationConfig {
            opt_level: "none".to_string(),
            enable_alias_analysis: false,
        };
        let result = compile_class_to_object_with_config(&class, &tmp, &config);
        assert!(result.is_ok(), "AOT compilation with no opts failed: {:?}", result);
        let _ = std::fs::remove_file(&tmp);
    }

    #[test]
    fn test_aot_compile_returns() {
        let tmp = std::env::temp_dir().join("jvmrs_aot_test_returns.o");
        for (name, bytecode) in [
            ("void", vec![0xb1]),                          // return
            ("int", vec![0x03, 0xac]),                      // iconst_0, ireturn
            ("long", vec![0x09, 0xad]),                     // lconst_0, lreturn
            ("float", vec![0x0b, 0xae]),                    // fconst_0, freturn
            ("double", vec![0x0e, 0xaf]),                   // dconst_0, dreturn
            ("ref", vec![0x01, 0xb0]),                      // aconst_null, areturn
        ] {
            let class = make_test_class(&bytecode);
            let result = compile_class_to_object(&class, &tmp);
            assert!(result.is_ok(), "AOT {} return failed: {:?}", name, result);
        }
        let _ = std::fs::remove_file(&tmp);
    }
}
