use crate::opcode::OpCode;
use crate::readwrite::*;
use std::collections::HashMap;

#[derive(Debug, Clone)]
pub enum Operand {
    U32(u32),
    I32(i32),
    I64(i64),
    F64(f64),
    Bool(bool),
    Label(String),
}

#[derive(Debug, Clone)]
pub enum Asm {
    Op(OpCode, Vec<Operand>),
    Label(String),
}

pub fn update_pc_by_operand(initial_pc: u32, operands: &Vec<Operand>) -> u32 {
    let mut pcp: u32 = initial_pc;
    for operand in operands {
        match operand {
            Operand::U32(_) | Operand::I32(_) | Operand::Label(_) => pcp += 4,
            Operand::I64(_) | Operand::F64(_) => pcp += 8,
            Operand::Bool(_) => pcp += 1,
        }
    }
    pcp
}

pub fn assemble(instructions: &[Asm]) -> Result<Vec<u8>, String> {
    let mut labels = HashMap::new();
    let mut current_pc = 0u32;

    // First pass: resolve labels
    for instr in instructions {
        match instr {
            Asm::Label(name) => {
                labels.insert(name.clone(), current_pc);
            }
            Asm::Op(_opcode, operands) => {
                current_pc += 2; // OpCode is u16 (2 bytes)
                current_pc = update_pc_by_operand(current_pc, operands);
            }
        }
    }

    // Second pass: generate bytes
    let mut rom = Vec::new();
    for instr in instructions {
        match instr {
            Asm::Label(_) => {}
            Asm::Op(opcode, operands) => {
                write_u16_le(&mut rom, *opcode as u16).map_err(|e| e.to_string())?;
                for operand in operands {
                    match operand {
                        Operand::U32(v) => write_u32_le(&mut rom, *v).map_err(|e| e.to_string())?,
                        Operand::I32(v) => write_u32_le(&mut rom, *v as u32).map_err(|e| e.to_string())?,
                        Operand::I64(v) => write_i64_le(&mut rom, *v).map_err(|e| e.to_string())?,
                        Operand::F64(v) => write_f64_le(&mut rom, *v).map_err(|e| e.to_string())?,
                        Operand::Bool(v) => rom.push(if *v { 1 } else { 0 }),
                        Operand::Label(name) => {
                            let addr = labels.get(name).ok_or(format!("Undefined label: {}", name))?;
                            write_u32_le(&mut rom, *addr).map_err(|e| e.to_string())?;
                        }
                    }
                }
            }
        }
    }

    Ok(rom)
}