//! Assembly text output formatter for Xbox 360 disassembly.

use std::collections::HashMap;
use std::io::Write;

use xenia_xex::header::ImportLibrary;
use xenia_xex::pe::PeSection;

use crate::disasm::enrich_section;
use crate::func::FuncAnalysis;
use crate::sinks::text::write_instr_line;
use crate::xref::{XrefKind, Xref, XrefMap, resolve_source_label};

/// Metadata passed to the formatter (avoids exposing full Xex2Header internals).
pub struct DisasmInfo<'a> {
    pub image_base: u32,
    pub entry_point: u32,
    pub original_pe_name: Option<&'a str>,
    pub title_id: Option<u32>,
    pub media_id: Option<u32>,
    pub sections: &'a [PeSection],
    pub import_libraries: &'a [ImportLibrary],
}

/// Write full disassembly to the output stream.
pub fn write_asm(
    out: &mut dyn Write,
    pe: &[u8],
    info: &DisasmInfo,
    func_analysis: &FuncAnalysis,
    labels: &HashMap<u32, String>,
    import_map: &HashMap<u32, String>,
    xrefs: &XrefMap,
    data_annotations: &HashMap<u32, (u32, XrefKind)>,
) -> anyhow::Result<()> {
    // Header
    writeln!(out, "; ============================================================================")?;
    writeln!(out, "; Xbox 360 Disassembly — generated by xenia-rs")?;
    if let Some(name) = info.original_pe_name {
        writeln!(out, "; Original PE: {name}")?;
    }
    if let (Some(title_id), Some(media_id)) = (info.title_id, info.media_id) {
        writeln!(out, "; Title ID: 0x{title_id:08X}  Media ID: 0x{media_id:08X}")?;
    }
    writeln!(out, "; Image base: 0x{:08X}  Entry point: 0x{:08X}", info.image_base, info.entry_point)?;
    writeln!(out, "; Functions detected: {}", func_analysis.functions.len())?;
    writeln!(out, "; ============================================================================")?;
    writeln!(out)?;

    // Import declarations
    if !info.import_libraries.is_empty() {
        writeln!(out, "; ── Imports ─────────────────────────────────────────────────────────────────")?;
        for lib in info.import_libraries {
            writeln!(out, "; Library: {}", lib.name)?;
            for imp in &lib.imports {
                let resolved = crate::resolve_ordinal(&lib.name, imp.ordinal);
                let name = resolved.unwrap_or("???");
                let kind = if imp.record_type == 1 { "thunk" } else { "var" };
                writeln!(out, ";   [{kind}] 0x{:08X}  ordinal 0x{:04X} = {}", imp.address, imp.ordinal, name)?;
            }
        }
        writeln!(out)?;
    }

    // Disassemble each section
    for section in info.sections {
        writeln!(out, "; ── Section: {:8} VA=0x{:08X} Size=0x{:08X} Flags=0x{:08X} ──",
            section.name, section.virtual_address, section.virtual_size, section.flags)?;

        let va_start = section.virtual_address;
        let va_end = va_start + section.virtual_size;
        let file_start = section.virtual_address as usize;

        // Pre-sort data labels in this section for break-at-label hex dump
        let section_labels_sorted: Vec<u32> = if !section.is_code() {
            let sec_start = info.image_base + va_start;
            let sec_end = info.image_base + va_end;
            let mut addrs: Vec<u32> = labels.keys()
                .filter(|&&a| a >= sec_start && a < sec_end)
                .copied()
                .collect();
            addrs.sort();
            addrs
        } else {
            Vec::new()
        };

        if section.is_code() {
            writeln!(out, ".text")?;
            writeln!(out)?;

            let mut in_function = false;
            let abs_start = info.image_base + va_start;
            let abs_end = info.image_base + va_end;

            let items = enrich_section(
                pe, info.image_base, &section.name, abs_start, abs_end, func_analysis, labels,
            );
            for ri in items {
                let abs_addr = ri.item.addr;

                // Function start? Emit separator + header
                if let Some(fi) = func_analysis.get(abs_addr) {
                    if in_function {
                        writeln!(out, "; end function")?;
                    }
                    writeln!(out)?;
                    writeln!(out, "; ──────────────────────────────────────────────────────────────────────────")?;

                    let lbl = labels.get(&abs_addr).cloned()
                        .unwrap_or_else(|| format!("sub_{abs_addr:08X}"));

                    if fi.is_saverestore {
                        writeln!(out, "; FUNCTION: {lbl}  (save/restore GPR helper)")?;
                    } else if fi.is_leaf {
                        writeln!(out, "; FUNCTION: {lbl}  (leaf)")?;
                    } else {
                        let mut details = Vec::new();
                        if fi.frame_size > 0 {
                            details.push(format!("frame={}", fi.frame_size));
                        }
                        if fi.saved_gprs > 0 {
                            let first_reg = 32 - fi.saved_gprs;
                            details.push(format!("saves r{first_reg}-r31"));
                        }
                        let detail_str = if details.is_empty() {
                            String::new()
                        } else {
                            format!("  ({})", details.join(", "))
                        };
                        writeln!(out, "; FUNCTION: {lbl}{detail_str}")?;
                    }

                    if let Some(xref_lines) = format_xrefs(abs_addr, xrefs, func_analysis, labels) {
                        for line in &xref_lines {
                            writeln!(out, "{line}")?;
                        }
                    }

                    writeln!(out, "; ──────────────────────────────────────────────────────────────────────────")?;
                    in_function = true;
                }

                // Label
                if let Some(lbl) = labels.get(&abs_addr) {
                    if !func_analysis.is_function_start(abs_addr) {
                        writeln!(out)?;
                        if let Some(xref_lines) = format_xrefs(abs_addr, xrefs, func_analysis, labels) {
                            for line in &xref_lines {
                                writeln!(out, "{line}")?;
                            }
                        }
                        writeln!(out, "{lbl}:")?;
                    } else {
                        writeln!(out)?;
                        writeln!(out, "{lbl}:")?;
                    }
                }

                // Import thunk annotation
                if let Some(imp_name) = import_map.get(&abs_addr) {
                    writeln!(out, "        ; IMPORT: {imp_name}")?;
                }

                let data_annot = data_annotations.get(&abs_addr).copied();
                write_instr_line(out, &ri, labels, info.sections, info.image_base, data_annot)?;
            }
            if in_function {
                writeln!(out, "; end function")?;
            }
        } else {
            // Data section: hex dump
            writeln!(out, ".data")?;
            writeln!(out)?;

            let mut addr = va_start;
            while addr < va_end {
                let abs_addr = info.image_base + addr;
                let off = (addr - va_start) as usize + file_start;

                if let Some(lbl) = labels.get(&abs_addr) {
                    writeln!(out)?;
                    // Xrefs for data labels
                    if let Some(xref_lines) = format_xrefs(abs_addr, xrefs, func_analysis, labels) {
                        for line in &xref_lines {
                            writeln!(out, "{line}")?;
                        }
                    }
                    writeln!(out, "{lbl}:")?;
                }

                // Emit up to 16 bytes per line, but break at label boundaries
                let mut line_end = std::cmp::min(addr + 16, va_end);
                for &lbl_addr in &section_labels_sorted {
                    let lbl_va = lbl_addr - info.image_base;
                    if lbl_va > addr && lbl_va < line_end {
                        line_end = lbl_va;
                        break;
                    }
                    if lbl_va >= line_end { break; }
                }
                let byte_count = (line_end - addr) as usize;
                if off + byte_count > pe.len() { break; }

                write!(out, "  {:08X}:  ", abs_addr)?;
                for i in 0..byte_count {
                    write!(out, "{:02X}", pe[off + i])?;
                    if i % 4 == 3 { write!(out, " ")?; }
                }
                // ASCII representation
                let pad = (16 - byte_count) * 2 + (16 - byte_count) / 4;
                write!(out, "{:>width$}  |", "", width = pad)?;
                for i in 0..byte_count {
                    let b = pe[off + i];
                    let ch = if b.is_ascii_graphic() || b == b' ' { b as char } else { '.' };
                    write!(out, "{ch}")?;
                }
                writeln!(out, "|")?;

                addr = line_end;
            }
        }
        writeln!(out)?;
    }

    Ok(())
}

const XREF_DISPLAY_LIMIT: usize = 8;

fn format_xrefs(
    target: u32,
    xrefs: &XrefMap,
    func_analysis: &FuncAnalysis,
    labels: &HashMap<u32, String>,
) -> Option<Vec<String>> {
    let refs = xrefs.get(&target)?;
    if refs.is_empty() { return None; }

    let mut sorted: Vec<Xref> = refs.clone();
    sorted.sort();
    sorted.dedup();

    let total = sorted.len();
    let mut lines = Vec::new();

    let calls = sorted.iter().filter(|x| x.kind == XrefKind::Call).count();
    let jumps = sorted.iter().filter(|x| x.kind == XrefKind::Jump).count();
    let branches = sorted.iter().filter(|x| x.kind == XrefKind::Branch).count();
    let reads = sorted.iter().filter(|x| x.kind == XrefKind::DataRead).count();
    let writes = sorted.iter().filter(|x| x.kind == XrefKind::DataWrite).count();
    let data_refs = sorted.iter().filter(|x| x.kind == XrefKind::DataRef).count();

    let mut summary_parts = Vec::new();
    if calls > 0 { summary_parts.push(format!("{calls} call{}", if calls != 1 { "s" } else { "" })); }
    if jumps > 0 { summary_parts.push(format!("{jumps} jump{}", if jumps != 1 { "s" } else { "" })); }
    if branches > 0 { summary_parts.push(format!("{branches} branch{}", if branches != 1 { "es" } else { "" })); }
    if reads > 0 { summary_parts.push(format!("{reads} read{}", if reads != 1 { "s" } else { "" })); }
    if writes > 0 { summary_parts.push(format!("{writes} write{}", if writes != 1 { "s" } else { "" })); }
    if data_refs > 0 { summary_parts.push(format!("{data_refs} ref{}", if data_refs != 1 { "s" } else { "" })); }

    lines.push(format!("; XREF: {} ({})", summary_parts.join(", "), total));

    for (i, xref) in sorted.iter().enumerate() {
        if i >= XREF_DISPLAY_LIMIT {
            lines.push(format!(";   ... and {} more", total - XREF_DISPLAY_LIMIT));
            break;
        }
        let source_label = resolve_source_label(xref.source, func_analysis, labels);
        lines.push(format!(";   {} from {}", xref.kind.tag(), source_label));
    }

    Some(lines)
}