gaia-assembler 0.1.1

use crate::{
    backends::{Backend, GeneratedFiles},
    config::GaiaConfig,
    instruction::{CastKind, CmpCondition, CoreInstruction, GaiaInstruction, ManagedInstruction},
    program::{GaiaFunction, GaiaModule},
    types::GaiaType,
};
use gaia_binary::{BinaryWriter, Leb128};
use gaia_types::{
    helpers::{Architecture, ArtifactType, CompilationTarget},
    Result,
};
use nyar_assembler::program::{
    instructions::NyarCodec,
    pool::NyarConstantPool,
    types::{NyarChunk, NyarConstant, NyarModule},
    NyarInstruction,
};
use std::collections::HashMap;

/// Nyar backend for Gaia assembler.
///
/// This backend translates Gaia IR into Nyar VM instructions.
pub struct NyarBackend;

impl Backend for NyarBackend {
    fn name(&self) -> &'static str {
        "nyar"
    }

    fn primary_target(&self) -> CompilationTarget {
        CompilationTarget::new("nyar", "nyar", "nyar")
    }

    fn artifact_type(&self) -> ArtifactType {
        ArtifactType::Executable
    }

    fn match_score(&self, target: &CompilationTarget) -> f32 {
        if target.build == Architecture::Nyar {
            100.0
        }
        else {
            0.0
        }
    }

    fn generate(&self, program: &GaiaModule, config: &GaiaConfig) -> Result<GeneratedFiles> {
        let mut nyar_module = NyarModule::default();
        let mut constant_pool = NyarConstantPool::new();

        for func in &program.functions {
            let chunk = self.translate_function(func, &mut constant_pool)?;
            nyar_module.chunks.push(chunk);
        }

        nyar_module.constants = constant_pool.pool;

        Ok(GeneratedFiles {
            artifact_type: ArtifactType::Executable,
            files: HashMap::new(),
            custom: Some(std::sync::Arc::new(nyar_module)),
            diagnostics: vec![],
        })
    }
}

impl NyarBackend {
    fn translate_function(&self, func: &GaiaFunction, pool: &mut NyarConstantPool) -> Result<NyarChunk> {
        let mut chunk = NyarChunk::default();
        let mut instructions = Vec::new();
        let mut label_positions = std::collections::HashMap::new();

        // First pass: collect label positions and translate instructions
        for block in &func.blocks {
            // Record label position
            label_positions.insert(block.label.clone(), instructions.len());

            for inst in &block.instructions {
                self.translate_instruction(inst, &mut instructions, pool)?;
            }
            // Handle terminator
            self.translate_terminator(&block.terminator, &mut instructions, pool)?;
        }

        // Second pass: fix up jump offsets
        let mut fixed_instructions = Vec::new();
        for (i, inst) in instructions.iter().enumerate() {
            match inst {
                NyarInstruction::Jump(_) => {
                    // For now, keep as 0, we'll implement proper offset calculation later
                    fixed_instructions.push(NyarInstruction::Jump(0));
                }
                NyarInstruction::JumpIfFalse(_) => {
                    fixed_instructions.push(NyarInstruction::JumpIfFalse(0));
                }
                NyarInstruction::JumpIfTrue(_) => {
                    fixed_instructions.push(NyarInstruction::JumpIfTrue(0));
                }
                _ => {
                    fixed_instructions.push(inst.clone());
                }
            }
        }

        let mut writer = BinaryWriter::<Vec<u8>, Leb128>::new(Vec::new());
        for inst in fixed_instructions {
            inst.encode(&mut writer).map_err(|e| gaia_types::GaiaError::custom_error(e.to_string()))?;
        }
        chunk.code = writer.into_inner();

        Ok(chunk)
    }


    fn translate_instruction(
        &self,
        inst: &GaiaInstruction,
        out: &mut Vec<NyarInstruction>,
        pool: &mut NyarConstantPool,
    ) -> Result<()> {
        match inst {
            GaiaInstruction::Core(core) => self.translate_core(core, out, pool),
            GaiaInstruction::Managed(managed) => self.translate_managed(managed, out, pool),
            GaiaInstruction::Domain(_) => Ok(()),
        }
    }

    fn translate_core(
        &self,
        inst: &CoreInstruction,
        out: &mut Vec<NyarInstruction>,
        pool: &mut NyarConstantPool,
    ) -> Result<()> {
        use crate::program::GaiaConstant;
        match inst {
            CoreInstruction::PushConstant(c) => match c {
                GaiaConstant::I64(v) => out.push(NyarInstruction::I64Const(*v)),
                GaiaConstant::F64(v) => out.push(NyarInstruction::F64Const(*v)),
                GaiaConstant::I32(v) => out.push(NyarInstruction::I32Const(*v)),
                GaiaConstant::F32(v) => out.push(NyarInstruction::F32Const(*v)),
                GaiaConstant::String(s) => {
                    let idx = pool.add(NyarConstant::String(s.clone()));
                    out.push(NyarInstruction::Push(idx));
                }
                other => {
                    // 警告：未处理的常量类型
                    eprintln!("Warning: Unhandled constant type: {:?}", other);
                }
            },
            CoreInstruction::Add(ty) => {
                match ty {
                    GaiaType::I32 => out.push(NyarInstruction::I32Add),
                    GaiaType::I64 => out.push(NyarInstruction::I64Add),
                    GaiaType::F32 => out.push(NyarInstruction::F32Add),
                    GaiaType::F64 => out.push(NyarInstruction::F64Add),
                    other => {
                        // 警告：未处理的类型
                        eprintln!("Warning: Unhandled type in Add: {:?}", other);
                    }
                }
            }
            CoreInstruction::Sub(ty) => {
                match ty {
                    GaiaType::I32 => out.push(NyarInstruction::I32Sub),
                    GaiaType::I64 => out.push(NyarInstruction::I64Sub),
                    GaiaType::F32 => out.push(NyarInstruction::F32Sub),
                    GaiaType::F64 => out.push(NyarInstruction::F64Sub),
                    other => {
                        eprintln!("Warning: Unhandled type in Sub: {:?}", other);
                    }
                }
            }
            CoreInstruction::Mul(ty) => {
                match ty {
                    GaiaType::I32 => out.push(NyarInstruction::I32Mul),
                    GaiaType::I64 => out.push(NyarInstruction::I64Mul),
                    GaiaType::F32 => out.push(NyarInstruction::F32Mul),
                    GaiaType::F64 => out.push(NyarInstruction::F64Mul),
                    other => {
                        eprintln!("Warning: Unhandled type in Mul: {:?}", other);
                    }
                }
            }
            CoreInstruction::Div(ty) => {
                match ty {
                    GaiaType::I32 => out.push(NyarInstruction::I32DivS),
                    GaiaType::I64 => out.push(NyarInstruction::I64DivS),
                    GaiaType::F32 => out.push(NyarInstruction::F32Div),
                    GaiaType::F64 => out.push(NyarInstruction::F64Div),
                    other => {
                        eprintln!("Warning: Unhandled type in Div: {:?}", other);
                    }
                }
            }
            CoreInstruction::Rem(ty) => {
                match ty {
                    GaiaType::I32 => out.push(NyarInstruction::I32RemS),
                    GaiaType::I64 => out.push(NyarInstruction::I64RemS),
                    other => {
                        eprintln!("Warning: Unhandled type in Rem: {:?}", other);
                    }
                }
            }
            CoreInstruction::And(ty) => {
                match ty {
                    GaiaType::I32 => out.push(NyarInstruction::I32And),
                    GaiaType::I64 => out.push(NyarInstruction::I64And),
                    other => {
                        eprintln!("Warning: Unhandled type in And: {:?}", other);
                    }
                }
            }
            CoreInstruction::Or(ty) => {
                match ty {
                    GaiaType::I32 => out.push(NyarInstruction::I32Or),
                    GaiaType::I64 => out.push(NyarInstruction::I64Or),
                    other => {
                        eprintln!("Warning: Unhandled type in Or: {:?}", other);
                    }
                }
            }
            CoreInstruction::Xor(ty) => {
                match ty {
                    GaiaType::I32 => out.push(NyarInstruction::I32Xor),
                    GaiaType::I64 => out.push(NyarInstruction::I64Xor),
                    other => {
                        eprintln!("Warning: Unhandled type in Xor: {:?}", other);
                    }
                }
            }
            CoreInstruction::Shl(ty) => {
                match ty {
                    GaiaType::I32 => out.push(NyarInstruction::I32Shl),
                    GaiaType::I64 => out.push(NyarInstruction::I64Shl),
                    other => {
                        eprintln!("Warning: Unhandled type in Shl: {:?}", other);
                    }
                }
            }
            CoreInstruction::Shr(ty) => {
                match ty {
                    GaiaType::I32 => out.push(NyarInstruction::I32ShrS),
                    GaiaType::I64 => out.push(NyarInstruction::I64ShrS),
                    other => {
                        eprintln!("Warning: Unhandled type in Shr: {:?}", other);
                    }
                }
            }
            CoreInstruction::Neg(ty) => {
                match ty {
                    GaiaType::I32 => out.push(NyarInstruction::I32Neg),
                    GaiaType::I64 => out.push(NyarInstruction::I64Neg),
                    GaiaType::F32 => out.push(NyarInstruction::F32Neg),
                    GaiaType::F64 => out.push(NyarInstruction::F64Neg),
                    other => {
                        eprintln!("Warning: Unhandled type in Neg: {:?}", other);
                    }
                }
            }
            CoreInstruction::Not(ty) => {
                match ty {
                    GaiaType::I32 => out.push(NyarInstruction::I32Not),
                    GaiaType::I64 => out.push(NyarInstruction::I64Not),
                    other => {
                        eprintln!("Warning: Unhandled type in Not: {:?}", other);
                    }
                }
            }
            CoreInstruction::Ret => out.push(NyarInstruction::Return),
            CoreInstruction::LoadLocal(idx, _) => out.push(NyarInstruction::LoadLocal(*idx as u8)),
            CoreInstruction::StoreLocal(idx, _) => out.push(NyarInstruction::StoreLocal(*idx as u8)),
            CoreInstruction::LoadArg(idx, _) => out.push(NyarInstruction::LoadLocal(*idx as u8)),
            CoreInstruction::StoreArg(idx, _) => out.push(NyarInstruction::StoreLocal(*idx as u8)),
            CoreInstruction::Pop => out.push(NyarInstruction::Pop),
            CoreInstruction::Dup => out.push(NyarInstruction::Dup(0)),
            CoreInstruction::Cmp(cond, ty) => {
                match ty {
                    GaiaType::I32 => match cond {
                        CmpCondition::Eq => out.push(NyarInstruction::I32Eq),
                        CmpCondition::Ne => out.push(NyarInstruction::I32Ne),
                        CmpCondition::Lt => out.push(NyarInstruction::I32LtS),
                        CmpCondition::Le => out.push(NyarInstruction::I32LeS),
                        CmpCondition::Gt => out.push(NyarInstruction::I32GtS),
                        CmpCondition::Ge => out.push(NyarInstruction::I32GeS),
                    },
                    GaiaType::I64 => match cond {
                        CmpCondition::Eq => out.push(NyarInstruction::I64Eq),
                        CmpCondition::Ne => out.push(NyarInstruction::I64Ne),
                        CmpCondition::Lt => out.push(NyarInstruction::I64LtS),
                        CmpCondition::Le => out.push(NyarInstruction::I64LeS),
                        CmpCondition::Gt => out.push(NyarInstruction::I64GtS),
                        CmpCondition::Ge => out.push(NyarInstruction::I64GeS),
                    },
                    GaiaType::F32 => match cond {
                        CmpCondition::Eq => out.push(NyarInstruction::F32Eq),
                        CmpCondition::Ne => out.push(NyarInstruction::F32Ne),
                        CmpCondition::Lt => out.push(NyarInstruction::F32Lt),
                        CmpCondition::Le => out.push(NyarInstruction::F32Le),
                        CmpCondition::Gt => out.push(NyarInstruction::F32Gt),
                        CmpCondition::Ge => out.push(NyarInstruction::F32Ge),
                    },
                    GaiaType::F64 => match cond {
                        CmpCondition::Eq => out.push(NyarInstruction::F64Eq),
                        CmpCondition::Ne => out.push(NyarInstruction::F64Ne),
                        CmpCondition::Lt => out.push(NyarInstruction::F64Lt),
                        CmpCondition::Le => out.push(NyarInstruction::F64Le),
                        CmpCondition::Gt => out.push(NyarInstruction::F64Gt),
                        CmpCondition::Ge => out.push(NyarInstruction::F64Ge),
                    },
                    other => {
                        eprintln!("Warning: Unhandled type in Cmp: {:?}", other);
                    }
                }
            }
            CoreInstruction::Br(label) => {
                // 暂时使用 0 作为偏移量，后续需要计算实际偏移量
                out.push(NyarInstruction::Jump(0));
            }
            CoreInstruction::BrTrue(_label) => {
                // 暂时使用 0 作为偏移量，后续需要计算实际偏移量
                out.push(NyarInstruction::JumpIfTrue(0));
            }
            CoreInstruction::BrFalse(_label) => {
                // 暂时使用 0 作为偏移量，后续需要计算实际偏移量
                out.push(NyarInstruction::JumpIfFalse(0));
            }
            CoreInstruction::Call(func_name, args_count) => {
                let idx = pool.add(NyarConstant::String(func_name.clone()));
                out.push(NyarInstruction::Call(idx, *args_count as u8));
            }
            CoreInstruction::CallIndirect(args_count) => {
                out.push(NyarInstruction::CallClosure(*args_count as u8));
            }
            CoreInstruction::New(type_name) => {
                let idx = pool.add(NyarConstant::String(type_name.clone()));
                out.push(NyarInstruction::NewObject(idx));
            }
            CoreInstruction::NewArray(_ty, _is_length_on_stack) => {
                // 暂时使用 0 作为长度，后续需要处理长度计算
                out.push(NyarInstruction::NewArray(0));
            }
            CoreInstruction::LoadField(_obj_type, field_name) => {
                let idx = pool.add(NyarConstant::String(field_name.clone()));
                out.push(NyarInstruction::GetField(idx));
            }
            CoreInstruction::StoreField(_obj_type, field_name) => {
                let idx = pool.add(NyarConstant::String(field_name.clone()));
                out.push(NyarInstruction::SetField(idx));
            }
            CoreInstruction::LoadElement(_ty) => {
                out.push(NyarInstruction::GetElement);
            }
            CoreInstruction::StoreElement(_ty) => {
                out.push(NyarInstruction::SetElement);
            }
            CoreInstruction::ArrayLength => {
                out.push(NyarInstruction::SizeOf);
            }
            CoreInstruction::ArrayPush => {
                out.push(NyarInstruction::PushElementRight);
            }
            CoreInstruction::StructNew(struct_name) => {
                let idx = pool.add(NyarConstant::String(struct_name.clone()));
                out.push(NyarInstruction::NewObject(idx));
            }
            CoreInstruction::StructGet { struct_name: _, field_index, is_signed: _ } => {
                // 暂时使用字段名作为索引，后续需要优化
                let idx = pool.add(NyarConstant::String(format!("field_{}", field_index)));
                out.push(NyarInstruction::GetField(idx));
            }
            CoreInstruction::StructSet { struct_name: _, field_index } => {
                // 暂时使用字段名作为索引，后续需要优化
                let idx = pool.add(NyarConstant::String(format!("field_{}", field_index)));
                out.push(NyarInstruction::SetField(idx));
            }
            CoreInstruction::ArrayNew(array_name) => {
                pool.add(NyarConstant::String(array_name.clone()));
                out.push(NyarInstruction::NewArray(0));
            }
            CoreInstruction::ArrayGet { array_name: _, is_signed: _ } => {
                out.push(NyarInstruction::GetElement);
            }
            CoreInstruction::ArraySet(_array_name) => {
                out.push(NyarInstruction::SetElement);
            }
            CoreInstruction::Cast { from, to, kind: _ } => {
                // 实现类型转换指令
                match (from, to) {
                    (GaiaType::I32, GaiaType::I64) => out.push(NyarInstruction::I32Extend64S),
                    (GaiaType::I32, GaiaType::F32) => out.push(NyarInstruction::I32ToF32S),
                    (GaiaType::I32, GaiaType::F64) => out.push(NyarInstruction::I32ToF64S),
                    (GaiaType::I64, GaiaType::I32) => out.push(NyarInstruction::I32Trunc64S),
                    (GaiaType::I64, GaiaType::F32) => out.push(NyarInstruction::I64ToF32S),
                    (GaiaType::I64, GaiaType::F64) => out.push(NyarInstruction::I64ToF64S),
                    (GaiaType::F32, GaiaType::I32) => out.push(NyarInstruction::F32ToI32S),
                    (GaiaType::F32, GaiaType::I64) => out.push(NyarInstruction::F32ToI64S),
                    (GaiaType::F32, GaiaType::F64) => out.push(NyarInstruction::F32ToF64),
                    (GaiaType::F64, GaiaType::I32) => out.push(NyarInstruction::F64ToI32S),
                    (GaiaType::F64, GaiaType::I64) => out.push(NyarInstruction::F64ToI64S),
                    (GaiaType::F64, GaiaType::F32) => out.push(NyarInstruction::F64ToF32),
                    _ => {
                        // 对于其他类型转换，使用通用的 Cast 指令
                        let idx = pool.add(NyarConstant::String(to.to_string()));
                        out.push(NyarInstruction::Cast(idx));
                    }
                }
            }
            CoreInstruction::Alloca(_ty, count) => {
                // 内存分配指令，在 Nyar 中使用局部变量模拟
                for _ in 0..*count {
                    out.push(NyarInstruction::PushNone);
                }
            }
            CoreInstruction::Load(_ty) => {
                // 加载指令，在 Nyar 中使用 GetElement 模拟
                out.push(NyarInstruction::GetElement);
            }
            CoreInstruction::Store(_ty) => {
                // 存储指令，在 Nyar 中使用 SetElement 模拟
                out.push(NyarInstruction::SetElement);
            }
            CoreInstruction::Gep { base_type: _, indices } => {
                //  gep 指令，在 Nyar 中使用 GetElement 模拟
                for _ in 0..indices.len() {
                    out.push(NyarInstruction::GetElement);
                }
            }

            CoreInstruction::Label(_label) => {
                // 标签，暂时不处理
            }
            CoreInstruction::Throw => {
                // 抛出异常指令，使用 Halt 作为替代
                out.push(NyarInstruction::Halt);
            }
        }
        Ok(())
    }

    fn translate_managed(
        &self,
        inst: &ManagedInstruction,
        out: &mut Vec<NyarInstruction>,
        pool: &mut NyarConstantPool,
    ) -> Result<()> {
        match inst {
            ManagedInstruction::Initiate(args) => out.push(NyarInstruction::Initiate(*args as u8)),
            ManagedInstruction::Finalize => out.push(NyarInstruction::Finalize),
            ManagedInstruction::CallMethod { target: _, method, signature, is_virtual, call_site_id: _ } => {
                let idx = pool.add(NyarConstant::String(method.clone()));
                if *is_virtual {
                    out.push(NyarInstruction::CallVirtual(idx, signature.params.len() as u8));
                } else {
                    out.push(NyarInstruction::InvokeMethod(idx, signature.params.len() as u8));
                }
            }
            ManagedInstruction::CallStatic { target, method, signature } => {
                let idx = pool.add(NyarConstant::String(format!("{}.{}", target, method)));
                out.push(NyarInstruction::Call(idx, signature.params.len() as u8));
            }
            ManagedInstruction::Box(_ty) => {
                // 装箱操作，暂时不处理
            }
            ManagedInstruction::Unbox(_ty) => {
                // 拆箱操作，暂时不处理
            }
            ManagedInstruction::InstanceOf(ty) => {
                let idx = pool.add(NyarConstant::String(ty.to_string()));
                out.push(NyarInstruction::InstanceOf(idx));
            }
            ManagedInstruction::CheckCast(ty) => {
                let idx = pool.add(NyarConstant::String(ty.to_string()));
                out.push(NyarInstruction::CheckCast(idx));
            }
        }
        Ok(())
    }

    fn translate_terminator(
        &self,
        term: &crate::program::GaiaTerminator,
        out: &mut Vec<NyarInstruction>,
        pool: &mut NyarConstantPool,
    ) -> Result<()> {
        use crate::program::GaiaTerminator::*;
        match term {
            Return => out.push(NyarInstruction::Return),
            Halt => out.push(NyarInstruction::Halt),
            Jump(_) => {
                out.push(NyarInstruction::Jump(0));
            }
            Branch { .. } => {
                out.push(NyarInstruction::JumpIfFalse(0));
            }
            Call { callee, args_count, .. } => {
                let idx = pool.add(NyarConstant::String(callee.clone()));
                out.push(NyarInstruction::Call(idx, *args_count as u8));
            }
        }
        Ok(())
    }
}