use std::collections::HashMap;
use std::fmt;
use common::formats::binary_dse::DseExecutable;
use common::isa::instructions::Opcode;
use common::prelude::{Instruction, Register};

#[derive(Debug)]
pub struct Disassembly {
    pub symbols: Vec<SymbolLine>,
    pub text_lines: Vec<String>,
    pub data_lines: Vec<String>,
}

#[derive(Debug)]
pub struct SymbolLine {
    pub name: String,
    pub address: u32,
    pub section: SectionKind,
    pub binding: BindingKind,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SectionKind { Data, Text, Abs, Undef }
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BindingKind { Local, Global, Weak }

impl SectionKind {
    fn from_raw(v: u8) -> Self {
        match v { 1 => Self::Data, 2 => Self::Text, 0 => Self::Abs, _ => Self::Undef }
    }
}
impl BindingKind {
    fn from_raw(v: u8) -> Self {
        match v { 1 => Self::Global, 2 => Self::Weak, _ => Self::Local }
    }
}

pub fn disassemble(exe: &DseExecutable) -> Disassembly {
    let symbols = build_symbol_lines(exe);

    let mut text_labels: HashMap<u32, String> = HashMap::new();
    let mut data_labels: HashMap<u32, String> = HashMap::new();
    for sym in &symbols {
        match sym.section {
            SectionKind::Text => { text_labels.insert(sym.address, sym.name.clone()); }
            SectionKind::Data => { data_labels.insert(sym.address, sym.name.clone()); }
            _ => {}
        }
    }
    // symbols can refer to either section by full resolved address (per your
    // v1 loader convention: data lives right after text), so a single combined
    // map is actually what address-resolution needs — keep both, but merge
    // for lookups against full 32-bit addresses.
    let mut all_labels = text_labels.clone();
    all_labels.extend(data_labels.clone());

    let text_lines = build_text_lines(&exe.text, &all_labels);
    let data_lines = build_data_lines(&exe.data, &data_labels);

    Disassembly { symbols, text_lines, data_lines }
}

fn build_symbol_lines(exe: &DseExecutable) -> Vec<SymbolLine> {
    exe.symbols
        .iter()
        .map(|s| SymbolLine {
            name: read_string(&exe.string_table, s.name_offset),
            address: s.value,
            section: SectionKind::from_raw(s.section),
            binding: BindingKind::from_raw(s.binding),
        })
        .collect()
}

fn read_string(strtab: &[u8], offset: u32) -> String {
    let start = offset as usize;
    let end = strtab[start..].iter().position(|&b| b == 0).map(|p| start + p).unwrap_or(strtab.len());
    String::from_utf8_lossy(&strtab[start..end]).into_owned()
}

/// Address a jump-like instruction actually targets, if statically known.
/// Returns None when the jump is register-relative (addr reg != Zero) since
/// that depends on runtime state and can't be resolved here.
fn static_jump_target(addr_reg: Register, imm: u16) -> Option<u32> {
    (addr_reg == Register::Zero).then_some(imm as u32)
}

fn build_text_lines(words: &[u32], labels: &HashMap<u32, String>) -> Vec<String> {
    let mut out = Vec::new();
    let mut i = 0;

    while i < words.len() {
        let addr = 4 * i as u32;
        if let Some(label) = labels.get(&addr) {
            out.push(format!("{label}:"));
        }

        let instr = Instruction(words[i]);
        let opcode = Opcode::from_u8(instr.opcode());

        // Collapse Lli+Lui address-load pairs back into `lwi` (or a full
        // ldx/stx pseudo-op if a matching mem instruction follows), mirroring
        // what expand_ldx/expand_stx/expand_lwi originally expanded from.
        if let Some(Opcode::Lli) = opcode {
            if let Some(collapsed) = try_collapse_addr_load(words, i, labels) {
                out.push(format!("    {}", collapsed.text));
                i += collapsed.consumed;
                continue;
            }
        }

        out.push(format!("    {}", render_instruction(instr, opcode, labels)));
        i += 1;
    }

    out
}

struct Collapsed {
    text: String,
    consumed: usize,
}

fn try_collapse_addr_load(
    words: &[u32],
    i: usize,
    labels: &HashMap<u32, String>,
) -> Option<Collapsed> {
    let lli = Instruction(words[i]);
    let lui = Instruction(*words.get(i + 1)?);

    if Opcode::from_u8(lui.opcode())? as u8 != Opcode::Lui as u8 {
        return None;
    }
    if lui.dest_checked() != lli.dest_checked() {
        return None;
    }

    let dest = lli.dest_checked();
    let full_addr = (lli.imm16() as u32) | ((lui.imm16() as u32) << 16);
    let addr_text = labels.get(&full_addr).cloned().unwrap_or_else(|| format!("{full_addr:#010x}"));

    // Check whether a mem op immediately follows, referencing `dest` —
    // that's the ldb/ldh/ldw/stb/sth/stw expansion pattern, not a bare lwi.
    if let Some(&mem_word) = words.get(i + 2) {
        let mem = Instruction(mem_word);
        if let Some(mem_op) = Opcode::from_u8(mem.opcode()) {
            let is_mem_op = matches!(
                mem_op,
                Opcode::Ldb | Opcode::Ldbs | Opcode::Ldh | Opcode::Ldhs | Opcode::Ldw
                    | Opcode::Stb | Opcode::Sth | Opcode::Stw
            );
            if is_mem_op && (mem.src1_checked() == dest || mem.dest_checked() == dest) {
                let mnemonic = format!("{mem_op:?}").to_lowercase();
                let other = if mem.src1_checked() == dest { mem.dest_checked() } else { mem.src1_checked() };
                return Some(Collapsed {
                    text: format!("{mnemonic} {addr_text}, {other:?}, {:#06x}", mem.imm16()),
                    consumed: 3,
                });
            }
        }
    }

    Some(Collapsed {
        text: format!("lwi {addr_text}, {dest:?}"),
        consumed: 2,
    })
}

fn render_instruction(instr: Instruction, opcode: Option<Opcode>, labels: &HashMap<u32, String>) -> String {
    let Some(op) = opcode else {
        return format!("; INVALID OPCODE {:#04x} (word {:#010x})", instr.opcode(), instr.0);
    };

    let mnemonic = format!("{op:?}").to_lowercase();

    match op {
        Opcode::Nop | Opcode::Hlt | Opcode::Ret | Opcode::IRet => mnemonic,

        Opcode::Push => format!("{mnemonic} {:?}", instr.src1_checked()),
        Opcode::Pop => format!("{mnemonic} {:?}", instr.dest_checked()),

        Opcode::Lli | Opcode::Lui => {
            format!("{mnemonic} {:?}, {:#06x}", instr.dest_checked(), instr.imm16())
        }

        Opcode::Int => format!("{mnemonic} {:#06x}", instr.imm16()),

        Opcode::Ldb | Opcode::Ldbs | Opcode::Ldh | Opcode::Ldhs | Opcode::Ldw
        | Opcode::Stb | Opcode::Sth | Opcode::Stw | Opcode::Addi | Opcode::Subi => format!(
            "{mnemonic} {:?}, {:?}, {:#06x}",
            instr.src1_checked(), instr.dest_checked(), instr.imm16()
        ),

        Opcode::Jez | Opcode::Jnz => {
            let addr_reg = instr.dest_checked();
            let target = static_jump_target(addr_reg, instr.imm16())
                .and_then(|a| labels.get(&a).cloned());
            match target {
                Some(label) => format!("{mnemonic} {:?}, {label}", instr.src1_checked()),
                None if addr_reg == Register::Zero => {
                    format!("{mnemonic} {:?}, {:#06x}", instr.src1_checked(), instr.imm16())
                }
                None => format!(
                    "{mnemonic} {:?}, {:#06x}, {addr_reg:?}",
                    instr.src1_checked(), instr.imm16()
                ),
            }
        }

        Opcode::Jmp | Opcode::Call | Opcode::Jic | Opcode::Jnc => {
            let addr_reg = instr.dest_checked();
            let target = static_jump_target(addr_reg, instr.imm16())
                .and_then(|a| labels.get(&a).cloned());
            match target {
                Some(label) => format!("{mnemonic} {label}"),
                None if addr_reg == Register::Zero => format!("{mnemonic} {:#06x}", instr.imm16()),
                None => format!("{mnemonic} {:#06x}, {addr_reg:?}", instr.imm16()),
            }
        }

        Opcode::Mov | Opcode::Not => {
            format!("{mnemonic} {:?}, {:?}", instr.src1_checked(), instr.dest_checked())
        }

        Opcode::CMov => format!(
            "{mnemonic} {:?}, {:?}, {:?}",
            instr.src1_checked(), instr.dest_checked(), instr.src2_checked()
        ),

        Opcode::Add | Opcode::Sub | Opcode::And | Opcode::Nand | Opcode::Or | Opcode::Nor
        | Opcode::Xor | Opcode::Xnor | Opcode::Ieq | Opcode::Ine | Opcode::Ilt | Opcode::Ile
        | Opcode::Igt | Opcode::Ige => format!(
            "{mnemonic} {:?}, {:?}, {:?}",
            instr.src1_checked(), instr.src2_checked(), instr.dest_checked()
        ),

        Opcode::Shl | Opcode::Shr => format!(
            "{mnemonic} {:?}, {:?}, {:?}, {}",
            instr.src1_checked(), instr.dest_checked(), instr.src2_checked(), instr.shamt()
        ),
    }
}

fn build_data_lines(data: &[u8], data_labels: &HashMap<u32, String>) -> Vec<String> {
    let mut out = Vec::new();
    let mut i = 0;

    while i < data.len() {
        let addr = i as u32;
        if let Some(label) = data_labels.get(&addr) {
            out.push(format!("{label}:"));
        }

        // Try string detection only when this byte could actually start one.
        if data[i] != 0 {
            if let Some(end) = data[i..].iter().position(|&b| b == 0) {
                let slice = &data[i..i + end];
                if slice.iter().all(|&b| b.is_ascii_graphic() || b == b' ') {
                    out.push(format!("    db \"{}\"", String::from_utf8_lossy(slice).escape_default()));
                    i += end + 1; // skip the null terminator too
                    continue;
                }
            }
        }

        // Not a string start: consume only the run of zero bytes here,
        // grouped into 4-byte words, stopping the instant a non-zero byte
        // appears so we never swallow the start of the next string/value.
        let run_end = data[i..]
            .iter()
            .position(|&b| b != 0)
            .map(|p| i + p)
            .unwrap_or(data.len());

        if run_end == i {
            // data[i] is non-zero but didn't form a printable/terminated
            // string (e.g. binary constant) — dump one raw word.
            let mut buf = [0u8; 4];
            let take = (data.len() - i).min(4);
            buf[..take].copy_from_slice(&data[i..i + take]);
            out.push(format!("    dw {:#010x}", u32::from_le_bytes(buf)));
            i += take;
        } else {
            let mut j = i;
            while j < run_end {
                let take = (run_end - j).min(4);
                let mut buf = [0u8; 4];
                buf[..take].copy_from_slice(&data[j..j + take]);
                out.push(format!("    dw {:#010x}", u32::from_le_bytes(buf)));
                j += take;
            }
            i = j;
        }
    }

    out
}

impl fmt::Display for SectionKind {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Data => write!(f, "DATA"),
            Self::Text => write!(f, "TEXT"),
            Self::Abs => write!(f, "ABS"),
            Self::Undef => write!(f, "UNDEF"),
        }
    }
}

impl fmt::Display for BindingKind {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Local => write!(f, "local"),
            Self::Global => write!(f, "global"),
            Self::Weak => write!(f, "weak"),
        }
    }
}

impl fmt::Display for SymbolLine {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "; {:<24} {:#010x}  {:<5} {}",
            self.name, self.address, self.section, self.binding
        )
    }
}

impl fmt::Display for Disassembly {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        writeln!(f, "; ==== symbols ====")?;
        for sym in &self.symbols {
            writeln!(f, "{sym}")?;
        }

        writeln!(f, "\n; ==== data ====")?;
        for line in &self.data_lines {
            writeln!(f, "{line}")?;
        }

        writeln!(f, "\n; ==== text ====")?;
        for line in &self.text_lines {
            writeln!(f, "{line}")?;
        }

        Ok(())
    }
}