use crate::instruction::{Instruction, Operand, MemoryOperand, Prefix};
use crate::parser::utilities::*;
use colored::*;

pub fn disassemble(data: &[u8]) -> Result<Vec<Instruction>, String> {
    let mut instructions = Vec::new();
    let mut offset = 0;

    while offset < data.len() {
        let (instruction, size) = decode_instruction(data, offset)?;
        instructions.push(instruction);
        offset += size;
    }

    Ok(instructions)
}

fn decode_instruction(data: &[u8], offset: usize) -> Result<(Instruction, usize), String> {
    let mut current_offset = offset;
    let mut prefix = None;
    let mut rex_prefix = 0;

    // Handle prefixes
    loop {
        match data.get(current_offset) {
            Some(0xF0) => prefix = Some(Prefix::Lock),
            Some(0xF2) => prefix = Some(Prefix::Repne),
            Some(0xF3) => prefix = Some(Prefix::Rep),
            Some(0x2E) => prefix = Some(Prefix::SegmentOverride("CS".to_string())),
            Some(0x36) => prefix = Some(Prefix::SegmentOverride("SS".to_string())),
            Some(0x3E) => prefix = Some(Prefix::SegmentOverride("DS".to_string())),
            Some(0x26) => prefix = Some(Prefix::SegmentOverride("ES".to_string())),
            Some(0x64) => prefix = Some(Prefix::SegmentOverride("FS".to_string())),
            Some(0x65) => prefix = Some(Prefix::SegmentOverride("GS".to_string())),
            Some(0x40..=0x4F) => rex_prefix = data[current_offset],
            _ => break,
        }
        current_offset += 1;
    }

    let opcode = *data.get(current_offset).ok_or("Unexpected end of data")?;
    current_offset += 1;

    let (mnemonic, operands, additional_bytes) = match opcode {
        0x00..=0x3F => decode_arithmetic_logic(opcode, data, current_offset, rex_prefix)?,
        0x40..=0x47 => ("INC".to_string(), vec![Operand::Register(get_register(opcode - 0x40, rex_prefix))], 0),
        0x48..=0x4F => ("DEC".to_string(), vec![Operand::Register(get_register(opcode - 0x48, rex_prefix))], 0),
        0x50..=0x57 => ("PUSH".to_string(), vec![Operand::Register(get_register(opcode - 0x50, rex_prefix))], 0),
        0x58..=0x5F => ("POP".to_string(), vec![Operand::Register(get_register(opcode - 0x58, rex_prefix))], 0),
        0x68 => {
            let imm = read_i32_le(data, current_offset).ok_or("Invalid immediate")?;
            ("PUSH".to_string(), vec![Operand::Immediate(imm as i64)], 4)
        },
        0x6A => {
            let imm = read_i8(data, current_offset).ok_or("Invalid immediate")?;
            ("PUSH".to_string(), vec![Operand::Immediate(imm as i64)], 1)
        },
        0x70..=0x7F => {
            let offset = read_i8(data, current_offset).ok_or("Invalid offset")?;
            (get_conditional_jump_mnemonic(opcode), vec![Operand::RelativeOffset(offset as i32)], 1)
        },
        0x81 => decode_immediate_arithmetic(data, current_offset, rex_prefix)?,
        0x83 => decode_immediate_arithmetic_small(data, current_offset, rex_prefix)?,
        0x88 | 0x89 => decode_mov(data, current_offset, rex_prefix)?,
        0x8A | 0x8B => decode_mov(data, current_offset, rex_prefix)?,
        0x90 => ("NOP".to_string(), vec![], 0),
        0xB0..=0xB7 => {
            let imm = read_u8(data, current_offset).ok_or("Invalid immediate")?;
            ("MOV".to_string(), vec![Operand::Register(get_register(opcode - 0xB0, rex_prefix)), Operand::Immediate(imm as i64)], 1)
        },
        0xB8..=0xBF => {
            let imm = read_u32_le(data, current_offset).ok_or("Invalid immediate")?;
            ("MOV".to_string(), vec![Operand::Register(get_register(opcode - 0xB8, rex_prefix)), Operand::Immediate(imm as i64)], 4)
        },
        0xC3 => ("RET".to_string(), vec![], 0),
        0xE8 => {
            let offset = read_i32_le(data, current_offset).ok_or("Invalid offset")?;
            ("CALL".to_string(), vec![Operand::RelativeOffset(offset)], 4)
        },
        0xEB => {
            let offset = read_i8(data, current_offset).ok_or("Invalid offset")?;
            ("JMP".to_string(), vec![Operand::RelativeOffset(offset as i32)], 1)
        },
        0xFF => decode_ff_group(data, current_offset, rex_prefix)?,
        _ => (format!("DB {:02X}", opcode), vec![], 0),  // Unknown opcode, treat as data byte
    };

    let instruction_size = current_offset - offset + additional_bytes;
    let bytes: Vec<u8> = data[offset..offset + instruction_size].to_vec();

    Ok((Instruction {
        opcode,
        mnemonic,
        operands,
        bytes,
        address: offset,
        prefix,
    }, instruction_size))
}

fn decode_arithmetic_logic(opcode: u8, data: &[u8], offset: usize, rex_prefix: u8) -> Result<(String, Vec<Operand>, usize), String> {
    let mnemonic = match opcode & 0x38 {
        0x00 => "ADD",
        0x08 => "OR",
        0x10 => "ADC",
        0x18 => "SBB",
        0x20 => "AND",
        0x28 => "SUB",
        0x30 => "XOR",
        0x38 => "CMP",
        _ => unreachable!(),
    };

    if let Some(modrm) = data.get(offset) {
        let (dest, src, size) = decode_modrm(data, offset, rex_prefix)?;
        Ok((mnemonic.to_string(), vec![dest, src], size))
    } else {
        Ok((mnemonic.to_string(), vec![], 0))
    }
}

fn decode_mov(data: &[u8], offset: usize, rex_prefix: u8) -> Result<(String, Vec<Operand>, usize), String> {
    if let Some(modrm) = data.get(offset) {
        let (dest, src, size) = decode_modrm(data, offset, rex_prefix)?;
        Ok(("MOV".to_string(), vec![dest, src], size))
    } else {
        Ok(("MOV".to_string(), vec![], 0))
    }
}

fn decode_modrm(data: &[u8], offset: usize, rex_prefix: u8) -> Result<(Operand, Operand, usize), String> {

    let modrm = read_u8(data, offset).ok_or("Invalid ModR/M byte")?;
    let mod_bits = (modrm >> 6) & 0x03;
    let reg = (modrm >> 3) & 0x07;
    let rm = modrm & 0x07;

    let (dest, src, size) = match mod_bits {
        0b00 => {
            if rm == 0b100 {
                // SIB byte
                let sib = read_u8(data, offset + 1).ok_or("Invalid SIB byte")?;
                let base = sib & 0x07;
                let index = (sib >> 3) & 0x07;
                let scale = 1 << ((sib >> 6) & 0x03);
                
                let memory = MemoryOperand {
                    base: Some(get_register(base, rex_prefix)),
                    index: Some(get_register(index, rex_prefix)),
                    scale: Some(scale),
                    displacement: None,
                };
                (Operand::Memory(memory), Operand::Register(get_register(reg, rex_prefix)), 2)
            } else if rm == 0b101 {
                // Displacement only
                let disp = read_i32_le(data, offset + 1).ok_or("Invalid displacement")?;
                let memory = MemoryOperand {
                    base: None,
                    index: None,
                    scale: None,
                    displacement: Some(disp),
                };
                (Operand::Memory(memory), Operand::Register(get_register(reg, rex_prefix)), 5)
            } else {
                // Register indirect
                let memory = MemoryOperand {
                    base: Some(get_register(rm, rex_prefix)),
                    index: None,
                    scale: None,
                    displacement: None,
                };
                (Operand::Memory(memory), Operand::Register(get_register(reg, rex_prefix)), 1)
            }
        },
        0b01 => {
            // 8-bit displacement
            let disp = read_i8(data, offset + 1).ok_or("Invalid displacement")? as i32;
            let memory = MemoryOperand {
                base: Some(get_register(rm, rex_prefix)),
                index: None,
                scale: None,
                displacement: Some(disp),
            };
            (Operand::Memory(memory), Operand::Register(get_register(reg, rex_prefix)), 2)
        },
        0b10 => {
            // 32-bit displacement
            let disp = read_i32_le(data, offset + 1).ok_or("Invalid displacement")?;
            let memory = MemoryOperand {
                base: Some(get_register(rm, rex_prefix)),
                index: None,
                scale: None,
                displacement: Some(disp),
            };
            (Operand::Memory(memory), Operand::Register(get_register(reg, rex_prefix)), 5)
        },
        0b11 => {
            // Register to register
            (Operand::Register(get_register(rm, rex_prefix)), Operand::Register(get_register(reg, rex_prefix)), 1)
        },
        _ => unreachable!(),
    };

    Ok((dest, src, size))
}

fn get_register(index: u8, rex_prefix: u8) -> String {
    let extended = (rex_prefix & 0b1000) != 0;
    match index {
        0 => if extended { "RAX".to_string() } else { "EAX".to_string() },
        1 => if extended { "RCX".to_string() } else { "ECX".to_string() },
        2 => if extended { "RDX".to_string() } else { "EDX".to_string() },
        3 => if extended { "RBX".to_string() } else { "EBX".to_string() },
        4 => if extended { "RSP".to_string() } else { "ESP".to_string() },
        5 => if extended { "RBP".to_string() } else { "EBP".to_string() },
        6 => if extended { "RSI".to_string() } else { "ESI".to_string() },
        7 => if extended { "RDI".to_string() } else { "EDI".to_string() },
        _ => unreachable!(),
    }
}

pub fn format_disassembly(instructions: &[Instruction]) -> String {
    instructions
        .iter()
        .map(|instr| format_instruction(instr))
        .collect::<Vec<String>>()
        .join("\n")
}
fn format_instruction(instr: &Instruction) -> String {
    let address = format!("{:04X}:", instr.address).yellow();
    let bytes = instr.bytes
        .iter()
        .map(|b| format!("{:02X}", b))
        .collect::<Vec<String>>()
        .join(" ")
        .bright_black();
    let mnemonic = instr.mnemonic.bright_blue();
    
    let operands = instr.operands
        .iter()
        .map(|op| match op {
            Operand::Register(reg) => reg.bright_cyan().to_string(),
            Operand::Immediate(imm) => format!("0x{:X}", imm).bright_green().to_string(),
            Operand::Memory(mem) => format_memory_operand(mem).bright_magenta().to_string(),
            Operand::RelativeOffset(offset) => format!("0x{:X}", offset).bright_green().to_string(),
        })
        .collect::<Vec<String>>()
        .join(", ");

    format!("{} {:20} {:6} {}", address, bytes, mnemonic, operands)
}

fn format_memory_operand(mem: &MemoryOperand) -> String {
    let mut parts = Vec::new();

    if let Some(base) = &mem.base {
        parts.push(base.bright_cyan().to_string());
    }

    if let Some(index) = &mem.index {
        if let Some(scale) = mem.scale {
            if scale != 1 {
                parts.push(format!("{}*{}", index.bright_cyan(), scale.to_string().bright_green()));
            } else {
                parts.push(index.bright_cyan().to_string());
            }
        } else {
            parts.push(index.bright_cyan().to_string());
        }
    }

    if let Some(disp) = mem.displacement {
        if disp != 0 || parts.is_empty() {
            if disp >= 0 {
                parts.push(format!("+0x{:X}", disp).bright_green().to_string());
            } else {
                parts.push(format!("-0x{:X}", disp.abs()).bright_green().to_string());
            }
        }
    }

    format!("[{}]", parts.join(" "))
}

fn get_conditional_jump_mnemonic(opcode: u8) -> String {
    match opcode {
        0x70 => "JO",
        0x71 => "JNO",
        0x72 => "JB",
        0x73 => "JAE",
        0x74 => "JE",
        0x75 => "JNE",
        0x76 => "JBE",
        0x77 => "JA",
        0x78 => "JS",
        0x79 => "JNS",
        0x7A => "JP",
        0x7B => "JNP",
        0x7C => "JL",
        0x7D => "JGE",
        0x7E => "JLE",
        0x7F => "JG",
        _ => unreachable!(),
    }.to_string()
}

fn decode_immediate_arithmetic(data: &[u8], offset: usize, rex_prefix: u8) -> Result<(String, Vec<Operand>, usize), String> {
    let modrm = read_u8(data, offset).ok_or("Invalid ModR/M byte")?;
    let reg = (modrm >> 3) & 0x07;
    let (dest, _, size) = decode_modrm(data, offset, rex_prefix)?;
    let imm = read_i32_le(data, offset + size).ok_or("Invalid immediate")?;
    
    let mnemonic = match reg {
        0 => "ADD",
        1 => "OR",
        2 => "ADC",
        3 => "SBB",
        4 => "AND",
        5 => "SUB",
        6 => "XOR",
        7 => "CMP",
        _ => unreachable!(),
    };

    Ok((mnemonic.to_string(), vec![dest, Operand::Immediate(imm as i64)], size + 4))
}

fn decode_immediate_arithmetic_small(data: &[u8], offset: usize, rex_prefix: u8) -> Result<(String, Vec<Operand>, usize), String> {
    let modrm = read_u8(data, offset).ok_or("Invalid ModR/M byte")?;
    let reg = (modrm >> 3) & 0x07;
    let (dest, _, size) = decode_modrm(data, offset, rex_prefix)?;
    let imm = read_i8(data, offset + size).ok_or("Invalid immediate")?;
    
    let mnemonic = match reg {
        0 => "ADD",
        1 => "OR",
        2 => "ADC",
        3 => "SBB",
        4 => "AND",
        5 => "SUB",
        6 => "XOR",
        7 => "CMP",
        _ => unreachable!(),
    };

    Ok((mnemonic.to_string(), vec![dest, Operand::Immediate(imm as i64)], size + 1))
}


fn decode_ff_group(data: &[u8], offset: usize, rex_prefix: u8) -> Result<(String, Vec<Operand>, usize), String> {
    let modrm = read_u8(data, offset).ok_or("Invalid ModR/M byte")?;
    let reg = (modrm >> 3) & 0x07;
    let (operand, _, size) = decode_modrm(data, offset, rex_prefix)?;
    
    let mnemonic = match reg {
        0 => "INC",
        1 => "DEC",
        2 => "CALL",
        3 => "CALL",
        4 => "JMP",
        5 => "JMP",
        6 => "PUSH",
        7 => "POP",
        _ => "INVALID",
    };

    Ok((mnemonic.to_string(), vec![operand], size))
}

Find any instructions that are in here that aren't  implemented/are bad/don't make sense.

If they aren't there. Implement them. Provide the full code when you're finished:

CB80: 0F AC DB 07 31 D6    OR     [ESP -0x29CEF825], EBP
CB86: C5                   DB C5
CB87: E9                   DB E9
CB88: 73 FE                JAE    0xFFFFFFFE
CB8A: 0C C5                OR     EBP, EAX
CB8C: C9                   DB C9
CB8D: FE                   DB FE
CB8E: F6                   DB F6
CB8F: 89 C5                MOV    EBP, EAX
CB91: 03 7C 24             ADD    [ESP +0x24], EDI
CB94: 28 31                SUB    [ECX], ESI
CB96: CB                   DB CB
CB97: 0F A4 C0 05 C5 F1    OR     [ESP -0xE3AFA40], ESP
CB9D: 72 D2                JB     0xFFFFFFD2
CB9F: 1E C5                SBB    EBP, EAX
CBA1: C9                   DB C9
CBA2: EB F0                JMP    0xFFFFFFF0
CBA4: 01 F7                ADD    EDI, ESI
CBA6: 21 DD                AND    EBP, EBX
CBA8: 31 CB                XOR    EBX, ECX
CBAA: 01 C7                ADD    EDI, EAX
CBAC: C5                   DB C5
CBAD: E9                   DB E9
CBAE: 72 F2                JB     0xFFFFFFF2
CBB0: 02 C5                ADD    EBP, EAX
CBB2: F9                   DB F9
CBB3: 6F                   DB 6F
CBB4: 44 24 40 0F          AND    [EAX +0xF], EAX
CBB8: AC                   DB AC
CBB9: C0                   DB C0
CBBA: 07 31                ADD    [ECX], ESI