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

/// Represents an operand for an instruction.
#[derive(Debug, Clone)]
pub enum Operand {
    /// 32-bit unsigned integer (e.g., indices, addresses).
    U32(u32),
    /// 32-bit signed integer.
    I32(i32),
    /// 64-bit signed integer.
    I64(i64),
    /// 64-bit floating point.
    F64(f64),
    /// Boolean (true/false).
    Bool(bool),
    /// A symbolic label that will be resolved to an absolute PC address.
    Label(String),
    /// A symbolic label that will be resolved to a PC address relative to another label.
    RelLabel(String, String),
}

/// Represents an assembly-level element (either an instruction or a label).
#[derive(Debug, Clone)]
pub enum Asm {
    /// An OpCode followed by its operands. The mnemonics represent the operation to be performed.
    Op(OpCode, Vec<Operand>),
    /// A named marker in the code (e.g., "start:").
    Label(String),
}

/// Internal helper to calculate the next PC based on the operands' sizes.
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(_) | Operand::RelLabel(_, _) => pcp += 4,
            Operand::I64(_) | Operand::F64(_) => pcp += 8,
            Operand::Bool(_) => pcp += 1,
        }
    }
    pcp
}

pub struct AssembleResult {
    pub code: Vec<u8>,
    pub unresolved_labels: HashMap<String, Vec<u32>>,
}

/// Converts a list of assembly instructions into raw ROM bytes.
/// 
/// The assembly process is done in two passes:
/// 1. **Label Resolution**: Iterates through all instructions to calculate the PC (Program Counter)
///    of each label and stores them in a map.
/// 2. **Code Generation**: Translates each OpCode and its operands (resolving labels using the map)
///    into the final binary format.
pub fn assemble(instructions: &[Asm]) -> Result<Vec<u8>, String> {
    let res = assemble_with_unresolved(instructions)?;
    if !res.unresolved_labels.is_empty() {
        let labels: Vec<_> = res.unresolved_labels.keys().cloned().collect();
        return Err(format!("Undefined labels: {:?}", labels));
    }
    Ok(res.code)
}

pub fn assemble_with_unresolved(instructions: &[Asm]) -> Result<AssembleResult, 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();
    let mut unresolved_labels: HashMap<String, Vec<u32>> = HashMap::new();
    let mut pc = 0u32;

    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())?;
                pc += 2;
                for operand in operands {
                    match operand {
                        Operand::U32(v) => {
                            write_u32_le(&mut rom, *v).map_err(|e| e.to_string())?;
                            pc += 4;
                        }
                        Operand::I32(v) => {
                            write_u32_le(&mut rom, *v as u32).map_err(|e| e.to_string())?;
                            pc += 4;
                        }
                        Operand::I64(v) => {
                            write_i64_le(&mut rom, *v).map_err(|e| e.to_string())?;
                            pc += 8;
                        }
                        Operand::F64(v) => {
                            write_f64_le(&mut rom, *v).map_err(|e| e.to_string())?;
                            pc += 8;
                        }
                        Operand::Bool(v) => {
                            rom.push(if *v { 1 } else { 0 });
                            pc += 1;
                        }
                        Operand::Label(name) => {
                            if let Some(addr) = labels.get(name) {
                                write_u32_le(&mut rom, *addr).map_err(|e| e.to_string())?;
                            } else {
                                unresolved_labels.entry(name.clone()).or_default().push(pc);
                                write_u32_le(&mut rom, 0).map_err(|e| e.to_string())?; // Placeholder
                            }
                            pc += 4;
                        }
                        Operand::RelLabel(name, base) => {
                            let addr = labels.get(name).ok_or(format!("Undefined label: {}", name))?;
                            let base_addr = labels.get(base).ok_or(format!("Undefined base label: {}", base))?;
                            let rel_addr = (*addr as i64) - (*base_addr as i64);
                            write_u32_le(&mut rom, rel_addr as u32).map_err(|e| e.to_string())?;
                            pc += 4;
                        }
                    }
                }
            }
        }
    }

    Ok(AssembleResult {
        code: rom,
        unresolved_labels,
    })
}