gaia-assembler 0.1.1

use crate::{
    instruction::{CmpCondition, CoreInstruction, GaiaInstruction, ManagedInstruction},
    program::{GaiaBlock, GaiaConstant, GaiaFunction, GaiaModule},
    types::GaiaType,
};
use gaia_types::{GaiaError, Result};
use std::collections::HashMap;
#[cfg(feature = "x86_64-assembler")]
use x86_64_assembler::instruction::{Instruction, Operand, Register};

/// Relocation kind
#[derive(Debug, Clone)]
pub enum RelocationKind {
    /// RIP-relative addressing (32-bit offset)
    RipRelative,
    /// Relative 32-bit jump/call (rel32)
    Relative32,
    /// Absolute 64-bit address (imm64)
    Absolute64,
}

/// Relocation entry
#[derive(Debug, Clone)]
pub struct Relocation {
    /// Offset position in the instruction stream (relative to the start of the sequence)
    pub instruction_index: usize,
    /// Target symbol name (e.g., function name, label name, or data section offset)
    pub target: String,
    /// Relocation kind
    pub kind: RelocationKind,
    /// Additional offset
    pub addend: i32,
}

/// x86_64 instruction emitter, converts Gaia IR to abstract instruction sequences and relocation info
#[cfg(feature = "x86_64-assembler")]
pub struct X64Emitter<'a> {
    program: &'a GaiaModule,
    instructions: Vec<Instruction>,
    relocations: Vec<Relocation>,
    string_table: HashMap<String, usize>,
    rdata_content: Vec<u8>,
}

#[cfg(feature = "x86_64-assembler")]
impl<'a> X64Emitter<'a> {
    pub fn new(program: &'a GaiaModule) -> Self {
        Self {
            program,
            instructions: Vec::new(),
            relocations: Vec::new(),
            string_table: HashMap::new(),
            rdata_content: Vec::new(),
        }
    }

    /// Perform emission logic
    pub fn emit(&mut self) -> Result<()> {
        self.collect_strings();

        // sub rsp, 40 (Entry stack alignment)
        self.push_inst(Instruction::Sub { dst: Operand::Reg(Register::RSP), src: Operand::Imm { value: 40, size: 8 } });

        // 1. Generate Entry Stub
        self.emit_entry_stub()?;

        // 2. Generate all function bodies
        for function in &self.program.functions {
            self.emit_function(function)?;
        }

        // add rsp, 40
        self.push_inst(Instruction::Add { dst: Operand::Reg(Register::RSP), src: Operand::Imm { value: 40, size: 8 } });

        Ok(())
    }

    /// Extract all string constants used in the module and store them in the data section
    fn collect_strings(&mut self) {
        let mut next_offset = 0;
        for function in &self.program.functions {
            for block in &function.blocks {
                for inst in &block.instructions {
                    if let Some(s) = self.get_string_constant(inst) {
                        if !self.string_table.contains_key(s) {
                            self.string_table.insert(s.clone(), next_offset);
                            self.rdata_content.extend_from_slice(s.as_bytes());
                            self.rdata_content.push(0); // Null terminator
                            next_offset += s.len() + 1;
                        }
                    }
                }
            }
        }
    }

    fn get_string_constant<'b>(&self, inst: &'b GaiaInstruction) -> Option<&'b String> {
        match inst {
            GaiaInstruction::Core(CoreInstruction::PushConstant(GaiaConstant::String(s)))
            | GaiaInstruction::Core(CoreInstruction::New(s))
            | GaiaInstruction::Core(CoreInstruction::StoreField(_, s))
            | GaiaInstruction::Core(CoreInstruction::LoadField(_, s))
            | GaiaInstruction::Managed(ManagedInstruction::CallMethod { method: s, .. }) => Some(s),
            _ => None,
        }
    }

    fn emit_entry_stub(&mut self) -> Result<()> {
        // 设置返回值为 0
        self.push_inst(Instruction::Xor { dst: Operand::Reg(Register::EAX), src: Operand::Reg(Register::EAX) });
        // 退出程序
        self.push_inst(Instruction::Mov { dst: Operand::Reg(Register::EAX), src: Operand::Imm { value: 0, size: 32 } });
        self.push_inst(Instruction::Ret);

        Ok(())
    }

    fn emit_function(&mut self, function: &GaiaFunction) -> Result<()> {
        // Record function label
        self.push_inst(Instruction::Label(function.name.clone()));

        // --- Prologue ---
        self.push_inst(Instruction::Push { op: Operand::Reg(Register::RBP) });
        self.push_inst(Instruction::Mov { dst: Operand::Reg(Register::RBP), src: Operand::Reg(Register::RSP) });

        // Calculate stack size (Locals + Shadow Space)
        let locals_count = function
            .blocks
            .iter()
            .flat_map(|b| &b.instructions)
            .filter(|i| matches!(i, GaiaInstruction::Core(CoreInstruction::Alloca(_, _))))
            .count();
        let has_managed_calls =
            function.blocks.iter().flat_map(|b| &b.instructions).any(|i| matches!(i, GaiaInstruction::Managed(_)));

        let locals_size = locals_count * 8;
        let shadow_space = if has_managed_calls { 64 } else { 32 };
        let total_stack_size = (locals_size + shadow_space + 15) & !15;

        if total_stack_size > 0 {
            self.push_inst(Instruction::Sub {
                dst: Operand::Reg(Register::RSP),
                src: Operand::Imm { value: total_stack_size as i64, size: 32 },
            });
        }

        // Save first 4 parameters into Shadow Space (Windows x64 calling convention)
        self.push_inst(Instruction::Mov {
            dst: Operand::Mem { base: Some(Register::RBP), index: None, scale: 1, displacement: 0x10 },
            src: Operand::Reg(Register::RCX),
        });
        self.push_inst(Instruction::Mov {
            dst: Operand::Mem { base: Some(Register::RBP), index: None, scale: 1, displacement: 0x18 },
            src: Operand::Reg(Register::RDX),
        });
        self.push_inst(Instruction::Mov {
            dst: Operand::Mem { base: Some(Register::RBP), index: None, scale: 1, displacement: 0x20 },
            src: Operand::Reg(Register::R8),
        });
        self.push_inst(Instruction::Mov {
            dst: Operand::Mem { base: Some(Register::RBP), index: None, scale: 1, displacement: 0x28 },
            src: Operand::Reg(Register::R9),
        });

        // --- Function Body ---
        for block in &function.blocks {
            self.emit_block(block, total_stack_size)?;
        }

        Ok(())
    }

    fn emit_block(&mut self, block: &GaiaBlock, total_stack_size: usize) -> Result<()> {
        self.push_inst(Instruction::Label(block.label.clone()));

        for inst in &block.instructions {
            match inst {
                GaiaInstruction::Core(core_inst) => self.emit_core_inst(core_inst, total_stack_size)?,
                GaiaInstruction::Managed(managed_inst) => self.emit_managed_inst(managed_inst)?,
                _ => return Err(GaiaError::custom_error(format!("Unsupported: {:?}", inst))),
            }
        }
        Ok(())
    }

    fn emit_core_inst(&mut self, inst: &CoreInstruction, total_stack_size: usize) -> Result<()> {
        match inst {
            CoreInstruction::PushConstant(constant) => match constant {
                GaiaConstant::I64(v) => {
                    self.push_inst(Instruction::Mov {
                        dst: Operand::Reg(Register::RAX),
                        src: Operand::Imm { value: *v, size: 64 },
                    });
                    self.push_inst(Instruction::Push { op: Operand::Reg(Register::RAX) });
                }
                GaiaConstant::I32(v) => {
                    self.push_inst(Instruction::Push { op: Operand::Imm { value: *v as i64, size: 32 } });
                }
                GaiaConstant::String(s) => {
                    let offset = *self.string_table.get(s).unwrap() as i32;
                    self.push_reloc(".rdata", RelocationKind::RipRelative, offset);
                    self.push_inst(Instruction::Lea { dst: Register::RAX, displacement: 0, rip_relative: true });
                    self.push_inst(Instruction::Push { op: Operand::Reg(Register::RAX) });
                }
                _ => {
                    // TODO: Other constant types
                }
            },
            CoreInstruction::Pop => {
                self.push_inst(Instruction::Pop { dst: Operand::Reg(Register::RAX) });
            }
            CoreInstruction::Add(_) => {
                self.push_inst(Instruction::Pop { dst: Operand::Reg(Register::RBX) });
                self.push_inst(Instruction::Pop { dst: Operand::Reg(Register::RAX) });
                self.push_inst(Instruction::Add { dst: Operand::Reg(Register::RAX), src: Operand::Reg(Register::RBX) });
                self.push_inst(Instruction::Push { op: Operand::Reg(Register::RAX) });
            }
            CoreInstruction::Sub(_) => {
                self.push_inst(Instruction::Pop { dst: Operand::Reg(Register::RBX) });
                self.push_inst(Instruction::Pop { dst: Operand::Reg(Register::RAX) });
                self.push_inst(Instruction::Sub { dst: Operand::Reg(Register::RAX), src: Operand::Reg(Register::RBX) });
                self.push_inst(Instruction::Push { op: Operand::Reg(Register::RAX) });
            }
            CoreInstruction::Mul(_) => {
                self.push_inst(Instruction::Pop { dst: Operand::Reg(Register::RBX) });
                self.push_inst(Instruction::Pop { dst: Operand::Reg(Register::RAX) });
                self.push_inst(Instruction::Imul { dst: Register::RAX, src: Operand::Reg(Register::RBX) });
                self.push_inst(Instruction::Push { op: Operand::Reg(Register::RAX) });
            }
            CoreInstruction::Div(_) => {
                self.push_inst(Instruction::Pop { dst: Operand::Reg(Register::RBX) });
                self.push_inst(Instruction::Pop { dst: Operand::Reg(Register::RAX) });
                self.push_inst(Instruction::Cqo);
                self.push_inst(Instruction::Idiv { src: Operand::Reg(Register::RBX) });
                self.push_inst(Instruction::Push { op: Operand::Reg(Register::RAX) });
            }
            CoreInstruction::Cmp(cond, _) => {
                self.push_inst(Instruction::Pop { dst: Operand::Reg(Register::RBX) });
                self.push_inst(Instruction::Pop { dst: Operand::Reg(Register::RAX) });
                self.push_inst(Instruction::Cmp { dst: Operand::Reg(Register::RAX), src: Operand::Reg(Register::RBX) });
                let cc = match cond {
                    CmpCondition::Eq => x86_64_assembler::instruction::Condition::E,
                    CmpCondition::Ne => x86_64_assembler::instruction::Condition::NE,
                    CmpCondition::Lt => x86_64_assembler::instruction::Condition::L,
                    CmpCondition::Le => x86_64_assembler::instruction::Condition::LE,
                    CmpCondition::Gt => x86_64_assembler::instruction::Condition::G,
                    CmpCondition::Ge => x86_64_assembler::instruction::Condition::GE,
                };
                self.push_inst(Instruction::Setcc { cond: cc, dst: Operand::Reg(Register::AL) });
                self.push_inst(Instruction::Movzx { dst: Register::RAX, src: Operand::Reg(Register::AL), size: 8 });
                self.push_inst(Instruction::Push { op: Operand::Reg(Register::RAX) });
            }
            CoreInstruction::LoadLocal(idx, _) => {
                let offset = -((*idx as i32 + 1) * 8);
                self.push_inst(Instruction::Mov {
                    dst: Operand::Reg(Register::RAX),
                    src: Operand::Mem { base: Some(Register::RBP), index: None, scale: 1, displacement: offset },
                });
                self.push_inst(Instruction::Push { op: Operand::Reg(Register::RAX) });
            }
            CoreInstruction::StoreLocal(idx, _) => {
                let offset = -((*idx as i32 + 1) * 8);
                self.push_inst(Instruction::Pop { dst: Operand::Reg(Register::RAX) });
                self.push_inst(Instruction::Mov {
                    dst: Operand::Mem { base: Some(Register::RBP), index: None, scale: 1, displacement: offset },
                    src: Operand::Reg(Register::RAX),
                });
            }
            CoreInstruction::LoadArg(idx, _) => {
                let offset = (*idx as i32 + 2) * 8;
                self.push_inst(Instruction::Mov {
                    dst: Operand::Reg(Register::RAX),
                    src: Operand::Mem { base: Some(Register::RBP), index: None, scale: 1, displacement: offset },
                });
                self.push_inst(Instruction::Push { op: Operand::Reg(Register::RAX) });
            }
            CoreInstruction::BrTrue(target) => {
                self.push_inst(Instruction::Pop { dst: Operand::Reg(Register::RAX) });
                self.push_inst(Instruction::Test { dst: Operand::Reg(Register::RAX), src: Operand::Reg(Register::RAX) });
                self.push_reloc(target, RelocationKind::Relative32, 0);
                self.push_inst(Instruction::Jcc {
                    cond: x86_64_assembler::instruction::Condition::NE,
                    target: Operand::Imm { value: 0, size: 32 },
                });
            }
            CoreInstruction::BrFalse(target) => {
                self.push_inst(Instruction::Pop { dst: Operand::Reg(Register::RAX) });
                self.push_inst(Instruction::Test { dst: Operand::Reg(Register::RAX), src: Operand::Reg(Register::RAX) });
                self.push_reloc(target, RelocationKind::Relative32, 0);
                self.push_inst(Instruction::Jcc {
                    cond: x86_64_assembler::instruction::Condition::E,
                    target: Operand::Imm { value: 0, size: 32 },
                });
            }
            CoreInstruction::Label(name) => {
                self.push_inst(Instruction::Label(name.clone()));
            }
            CoreInstruction::Ret => {
                self.push_inst(Instruction::Pop { dst: Operand::Reg(Register::RAX) });
                if total_stack_size > 0 {
                    self.push_inst(Instruction::Add {
                        dst: Operand::Reg(Register::RSP),
                        src: Operand::Imm { value: total_stack_size as i64, size: 32 },
                    });
                }
                self.push_inst(Instruction::Pop { dst: Operand::Reg(Register::RBP) });
                self.push_inst(Instruction::Ret);
            }
            CoreInstruction::Br(target) => {
                self.push_reloc(target, RelocationKind::Relative32, 0);
                self.push_inst(Instruction::Jmp { target: Operand::Imm { value: 0, size: 32 } });
            }
            CoreInstruction::Call(name, argc) => {
                // Setup arguments (first 4 in registers, rest on stack)
                self.emit_call_setup(*argc)?;

                // Determine if it's an internal or external call
                // Handled by relocation processor for now, which distinguishes symbol source
                self.push_reloc(name, RelocationKind::Relative32, 0);
                self.push_inst(Instruction::Call { target: Operand::Imm { value: 0, size: 32 } });
                self.push_inst(Instruction::Push { op: Operand::Reg(Register::RAX) });
            }
            CoreInstruction::New(_type_name) => {
                // Call gc_alloc from runtime
                // 1. Get type size (hardcoded for now or from metadata)
                self.push_inst(Instruction::Mov {
                    dst: Operand::Reg(Register::RCX),
                    src: Operand::Imm { value: 64, size: 64 }, // TODO: Dynamically calculate type size
                });
                self.push_reloc("gaia_gc_alloc", RelocationKind::RipRelative, 0);
                self.push_inst(Instruction::Call {
                    target: Operand::Mem { base: None, index: None, scale: 1, displacement: 0 },
                });
                self.push_inst(Instruction::Push { op: Operand::Reg(Register::RAX) });
            }
            CoreInstruction::LoadField(_type_name, _field_name) => {
                // Stack: [..., object_ptr]
                self.push_inst(Instruction::Pop { dst: Operand::Reg(Register::RAX) });
                // TODO: Get offset based on type_name and field_name
                let offset = 8; // Hardcoded for now
                self.push_inst(Instruction::Mov {
                    dst: Operand::Reg(Register::RAX),
                    src: Operand::Mem { base: Some(Register::RAX), index: None, scale: 1, displacement: offset },
                });
                self.push_inst(Instruction::Push { op: Operand::Reg(Register::RAX) });
            }
            CoreInstruction::StoreField(_type_name, _field_name) => {
                // Stack: [..., object_ptr, value]
                self.push_inst(Instruction::Pop { dst: Operand::Reg(Register::RBX) }); // value
                self.push_inst(Instruction::Pop { dst: Operand::Reg(Register::RAX) }); // object_ptr
                                                                                       // TODO: Get offset
                let offset = 8;
                self.push_inst(Instruction::Mov {
                    dst: Operand::Mem { base: Some(Register::RAX), index: None, scale: 1, displacement: offset },
                    src: Operand::Reg(Register::RBX),
                });
            }
            _ => { /* TODO: Complete other infrequent instructions */ }
        }
        Ok(())
    }

    fn emit_managed_inst(&mut self, inst: &ManagedInstruction) -> Result<()> {
        match inst {
            ManagedInstruction::CallMethod { method, .. } => {
                // Simplified: treat as ordinary Call
                self.push_reloc(method, RelocationKind::Relative32, 0);
                self.push_inst(Instruction::Call { target: Operand::Imm { value: 0, size: 32 } });
                self.push_inst(Instruction::Push { op: Operand::Reg(Register::RAX) });
            }
            ManagedInstruction::CallStatic { method, .. } => {
                self.push_reloc(method, RelocationKind::Relative32, 0);
                self.push_inst(Instruction::Call { target: Operand::Imm { value: 0, size: 32 } });
                self.push_inst(Instruction::Push { op: Operand::Reg(Register::RAX) });
            }
            _ => { /* TODO: Complete Managed instructions */ }
        }
        Ok(())
    }

    fn emit_call_setup(&mut self, argc: usize) -> Result<()> {
        if argc >= 4 {
            self.push_inst(Instruction::Pop { dst: Operand::Reg(Register::R9) });
        }
        if argc >= 3 {
            self.push_inst(Instruction::Pop { dst: Operand::Reg(Register::R8) });
        }
        if argc >= 2 {
            self.push_inst(Instruction::Pop { dst: Operand::Reg(Register::RDX) });
        }
        if argc >= 1 {
            self.push_inst(Instruction::Pop { dst: Operand::Reg(Register::RCX) });
        }
        Ok(())
    }

    // --- 辅助方法 ---

    fn push_inst(&mut self, inst: Instruction) {
        self.instructions.push(inst);
    }

    fn push_reloc(&mut self, target: &str, kind: RelocationKind, addend: i32) {
        self.relocations.push(Relocation {
            instruction_index: self.instructions.len(),
            target: target.to_string(),
            kind,
            addend,
        });
    }

    pub fn take_result(self) -> (Vec<Instruction>, Vec<Relocation>, Vec<u8>) {
        (self.instructions, self.relocations, self.rdata_content)
    }
}