// rasm-core/src/codegen/aot/coff.rs
//
//! # COFF64 Writer (`.obj`)
//!
//! Собирает объектный файл COFF64 для Windows x86_64 напрямую из
//! [`CompiledModule`]: каждая непустая секция модуля становится COFF-секцией,
//! символы — записями в Symbol Table, релокации — COFF-релокациями.
//!
//! Полученный `.obj` можно передать любому линкеру (link.exe, lld-link,
//! golink и т.д.) для создания финального исполняемого файла.
//!
//! ## Структура файла
//!
//! ```text
//! 0x00   20    COFF File Header
//! 0x14   40×N  Section Headers (по одной на каждую непустую секцию)
//! ...          Raw data каждой секции
//! ...          Relocations каждой секции (10 байт каждая)
//! ...          Symbol Table (18 байт каждая запись)
//! ...          String Table (4 байта длина + строки)
//! ```
use std::{
    collections::HashMap,
    io::{self, Write},
};

use crate::{
    codegen::reloc::RelocKind,
    object::{CompiledModule, SectionKind, SymbolVisibility},
};

// =========================================================================
// COFF Constants
// =========================================================================

const IMAGE_FILE_MACHINE_AMD64: u16 = 0x8664;

const IMAGE_SCN_CNT_CODE: u32 = 0x0000_0020;
const IMAGE_SCN_CNT_INITIALIZED_DATA: u32 = 0x0000_0040;
const IMAGE_SCN_CNT_UNINITIALIZED_DATA: u32 = 0x0000_0080;
const IMAGE_SCN_MEM_EXECUTE: u32 = 0x2000_0000;
const IMAGE_SCN_MEM_READ: u32 = 0x4000_0000;
const IMAGE_SCN_MEM_WRITE: u32 = 0x8000_0000;

const IMAGE_REL_AMD64_ADDR64: u16 = 0x0001;
const IMAGE_REL_AMD64_REL32: u16 = 0x0004;

const IMAGE_SYM_CLASS_EXTERNAL: u8 = 0x02;
const IMAGE_SYM_CLASS_STATIC: u8 = 0x03;

// =========================================================================
// COFFWriter
// =========================================================================

/// # COFF64 Writer
///
/// Принимает [`CompiledModule`] целиком и записывает готовый COFF64 `.obj`
/// в любой `Write`.
pub struct COFFWriter<'a> {
    module: &'a CompiledModule,
}

impl<'a> COFFWriter<'a> {
    pub fn new(module: &'a CompiledModule) -> Self {
        Self { module }
    }

    /// Записывает готовый COFF64-файл в переданный writer.
    pub fn write<W: Write>(&self, out: &mut W) -> io::Result<()> {
        let ctx = COFFContext::new(self.module);
        ctx.emit(out, self.module)
    }
}

// =========================================================================
// COFFContext
// =========================================================================

struct COFFContext {
    /// Секции (только непустые).
    sections: Vec<COFFSection>,
    /// Символы (сырые записи для Symbol Table).
    symbols: Vec<COFFSymbolRaw>,
    /// String Table байты.
    str_table: Vec<u8>,
    /// symbol name -> index in symbols.
    symbol_indices: HashMap<String, u32>,
}

#[derive(Debug, Clone)]
struct COFFSection {
    module_index: usize,
    name: [u8; 8],
    characteristics: u32,
    raw_data_offset: u32,
    raw_size: u32,
    reloc_offset: u32,
    num_relocs: u16,
}

#[derive(Debug, Clone)]
struct COFFSymbolRaw {
    /// Оригинальное имя (для поиска).
    name: String,
    value: u32,
    section_number: i16,
    typ: u16,
    storage_class: u8,
}

impl COFFContext {
    fn new(module: &CompiledModule) -> Self {
        let mut sections = Vec::new();
        let mut symbols = Vec::new();
        let mut symbol_indices = HashMap::new();

        // ---- 1. Секции ----
        for (idx, s) in module.sections.iter().enumerate() {
            if s.virtual_size == 0 {
                continue;
            }
            let is_bss = s.kind == SectionKind::Bss;
            sections.push(COFFSection {
                module_index: idx,
                name: coff_name_8(&s.name),
                characteristics: coff_characteristics(s.kind),
                raw_data_offset: 0,
                raw_size: if is_bss { 0 } else { s.bytes.len() as u32 },
                reloc_offset: 0,
                num_relocs: 0,
            });
        }

        // Функция: section_number по module.sections index (1-based).
        let sec_num_of = |mod_idx: usize| -> i16 {
            sections
                .iter()
                .position(|cs| cs.module_index == mod_idx)
                .map(|i| (i + 1) as i16)
                .unwrap_or(0)
        };

        // ---- 2. Определённые символы ----
        for sym in &module.symbols {
            if sym.visibility == SymbolVisibility::External && sym.section.is_none() {
                continue;
            }

            let sec_num = sym.section.map(|sid| sec_num_of(sid.0)).unwrap_or(0);
            let is_code = sym
                .section
                .and_then(|sid| module.section(sid))
                .map(|s| s.kind == SectionKind::Text)
                .unwrap_or(false);
            let typ: u16 = if is_code { 0x0020 } else { 0x0000 };
            let is_entry = module
                .entry
                .as_ref()
                .map(|e| e.symbol == sym.name)
                .unwrap_or(false);
            let sc = match sym.visibility {
                SymbolVisibility::Global => IMAGE_SYM_CLASS_EXTERNAL,
                _ if is_entry => IMAGE_SYM_CLASS_EXTERNAL,
                _ => IMAGE_SYM_CLASS_STATIC,
            };

            symbol_indices.insert(sym.name.clone(), symbols.len() as u32);
            symbols.push(COFFSymbolRaw {
                name: sym.name.clone(),
                value: sym.offset as u32,
                section_number: sec_num,
                typ,
                storage_class: sc,
            });
        }

        // ---- 3. Внешние символы (extern/import без определения) ----
        for sym in &module.symbols {
            if sym.visibility != SymbolVisibility::External || sym.section.is_some() {
                continue;
            }
            if symbol_indices.contains_key(&sym.name) {
                continue;
            }
            symbol_indices.insert(sym.name.clone(), symbols.len() as u32);
            symbols.push(COFFSymbolRaw {
                name: sym.name.clone(),
                value: 0,
                section_number: 0,
                typ: 0x0000,
                storage_class: IMAGE_SYM_CLASS_EXTERNAL,
            });
        }

        for import in &module.imports {
            if symbol_indices.contains_key(&import.alias) {
                continue;
            }
            symbol_indices.insert(import.alias.clone(), symbols.len() as u32);
            symbols.push(COFFSymbolRaw {
                name: import.alias.clone(),
                value: 0,
                section_number: 0,
                typ: 0x0000,
                storage_class: IMAGE_SYM_CLASS_EXTERNAL,
            });
        }

        // ---- 4. String Table ----
        let str_table = build_string_table(&symbols);

        Self {
            sections,
            symbols,
            str_table,
            symbol_indices,
        }
    }

    fn emit<W: Write>(&self, out: &mut W, module: &CompiledModule) -> io::Result<()> {
        // Вычисляем file offsets.
        let header_size: u32 = 20 + self.sections.len() as u32 * 40;
        let mut file_offset = header_size;

        let mut sections = self.sections.clone();
        for sec in &mut sections {
            sec.raw_data_offset = file_offset;

            let num: u16 = module
                .relocations
                .iter()
                .filter(|r| r.section.0 == sec.module_index)
                .count() as u16;
            sec.num_relocs = num;

            if sec.raw_size == 0 && num == 0 {
                continue;
            }

            if sec.raw_size > 0 {
                file_offset += sec.raw_size;
            }
            if num > 0 {
                sec.reloc_offset = file_offset;
                file_offset += num as u32 * 10;
            }
        }

        let symtab_offset = file_offset;
        let num_symbols = self.symbols.len() as u32;
        file_offset += num_symbols * 18;
        let _strtab_offset = file_offset; // not needed for writing

        // 1. File Header
        out.write_all(&file_header(
            sections.len() as u16,
            symtab_offset,
            num_symbols,
        ))?;

        // 2. Section Headers
        for sec in &sections {
            out.write_all(&section_header_bytes(sec))?;
        }

        // 3. Raw data + relocs
        for sec in &sections {
            if sec.raw_size > 0 {
                let mut raw = module.sections[sec.module_index].bytes.clone();
                for reloc in &module.relocations {
                    if reloc.section.0 != sec.module_index {
                        continue;
                    }
                    if reloc.kind == RelocKind::Rel32 {
                        let off = reloc.offset;
                        if off + 4 <= raw.len() {
                            raw[off..off + 4].copy_from_slice(&[0u8; 4]);
                        }
                    }
                }
                out.write_all(&raw)?;
            }
            if sec.num_relocs > 0 {
                for reloc in &module.relocations {
                    if reloc.section.0 != sec.module_index {
                        continue;
                    }
                    out.write_all(&self.reloc_bytes(reloc))?;
                }
            }
        }

        // 4. Symbol Table
        for sym in &self.symbols {
            out.write_all(&self.sym_entry_bytes(sym))?;
        }

        // 5. String Table
        out.write_all(&self.str_table)?;

        Ok(())
    }

    fn reloc_bytes(&self, reloc: &crate::object::ObjectRelocation) -> [u8; 10] {
        let mut r = [0u8; 10];
        put_u32(&mut r, 0, reloc.offset as u32);
        let sym_idx = self.symbol_indices.get(&reloc.target).copied().unwrap_or(0);
        put_u32(&mut r, 4, sym_idx);
        let typ = match reloc.kind {
            RelocKind::Abs64 => IMAGE_REL_AMD64_ADDR64,
            RelocKind::Rel32 => IMAGE_REL_AMD64_REL32,
        };
        put_u16(&mut r, 8, typ);
        r
    }

    fn sym_entry_bytes(&self, sym: &COFFSymbolRaw) -> [u8; 18] {
        let mut e = [0u8; 18];
        let name_bytes = sym_name_bytes(&sym.name, &self.str_table);
        e[0..8].copy_from_slice(&name_bytes);
        put_u32(&mut e, 8, sym.value);
        put_i16(&mut e, 12, sym.section_number);
        put_u16(&mut e, 14, sym.typ);
        e[16] = sym.storage_class;
        e
    }
}

// =========================================================================
// Вспомогательные функции
// =========================================================================

fn coff_name_8(name: &str) -> [u8; 8] {
    let mut buf = [0u8; 8];
    let bytes = name.as_bytes();
    let len = bytes.len().min(8);
    buf[..len].copy_from_slice(&bytes[..len]);
    buf
}

/// Возвращает 8-байтовое поле имени для Symbol Table Entry.
///
/// Если имя длиннее 8 символов, первые 4 байта = 0, а следующие 4 байта =
/// смещение этого имени (включая null-terminator) в String Table.
fn sym_name_bytes(name: &str, str_table: &[u8]) -> [u8; 8] {
    let mut buf = [0u8; 8];
    let bytes = name.as_bytes();
    if bytes.len() <= 8 {
        buf[..bytes.len()].copy_from_slice(bytes);
    } else {
        // Ищем offset в string table (первые 4 байта — общий размер,
        // затем идут null-terminated строки).
        let offset = find_in_str_table(name, str_table);
        buf[0..4].copy_from_slice(&0u32.to_le_bytes()); // zero = long name
        buf[4..8].copy_from_slice(&offset.to_le_bytes());
    }
    buf
}

/// Ищет null-terminated строку в String Table (начиная с offset 4).
fn find_in_str_table(name: &str, table: &[u8]) -> u32 {
    let mut pos: usize = 4; // пропускаем заголовок размера
    while pos < table.len() {
        // Ищем null-terminator
        let end = match table[pos..].iter().position(|&b| b == 0) {
            Some(n) => pos + n,
            None => break,
        };
        let entry = &table[pos..end];
        if entry == name.as_bytes() {
            return pos as u32;
        }
        pos = end + 1; // пропускаем null
    }
    0
}

/// Строит String Table (первые 4 байта = общий размер, затем
/// null-terminated строки для всех имён длиннее 8 символов).
fn build_string_table(symbols: &[COFFSymbolRaw]) -> Vec<u8> {
    let mut table = vec![0u8; 4]; // placeholder для размера

    for sym in symbols {
        if sym.name.len() > 8 && !string_already_in_table(&sym.name, &table) {
            table.extend_from_slice(sym.name.as_bytes());
            table.push(0);
            // patch offset into symbol name_bytes не делаем здесь —
            // sym_name_bytes() сделает это на лету при записи.
        }
    }

    let total = table.len() as u32;
    table[0..4].copy_from_slice(&total.to_le_bytes());
    table
}

fn string_already_in_table(name: &str, table: &[u8]) -> bool {
    if table.len() <= 4 {
        return false;
    }
    let mut pos = 4;
    while pos < table.len() {
        let end = match table[pos..].iter().position(|&b| b == 0) {
            Some(n) => pos + n,
            None => break,
        };
        if &table[pos..end] == name.as_bytes() {
            return true;
        }
        pos = end + 1;
    }
    false
}

// =========================================================================
// COFF бинарные структуры
// =========================================================================

fn file_header(num_sections: u16, symtab_offset: u32, num_symbols: u32) -> [u8; 20] {
    let mut h = [0u8; 20];
    put_u16(&mut h, 0, IMAGE_FILE_MACHINE_AMD64);
    put_u16(&mut h, 2, num_sections);
    put_u32(&mut h, 8, symtab_offset);
    put_u32(&mut h, 12, num_symbols);
    h
}

fn section_header_bytes(sec: &COFFSection) -> [u8; 40] {
    let mut h = [0u8; 40];
    h[0..8].copy_from_slice(&sec.name);
    put_u32(&mut h, 16, sec.raw_size);
    put_u32(&mut h, 20, sec.raw_data_offset);
    put_u32(&mut h, 24, sec.reloc_offset);
    put_u16(&mut h, 32, sec.num_relocs);
    put_u32(&mut h, 36, sec.characteristics);
    h
}

fn coff_characteristics(kind: SectionKind) -> u32 {
    match kind {
        SectionKind::Text => IMAGE_SCN_CNT_CODE | IMAGE_SCN_MEM_EXECUTE | IMAGE_SCN_MEM_READ,
        SectionKind::Data | SectionKind::Custom => {
            IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE
        }
        SectionKind::Bss => {
            IMAGE_SCN_CNT_UNINITIALIZED_DATA | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE
        }
    }
}

// =========================================================================
// Primitive writers
// =========================================================================

fn put_u16(b: &mut [u8], o: usize, v: u16) {
    b[o..o + 2].copy_from_slice(&v.to_le_bytes());
}

fn put_i16(b: &mut [u8], o: usize, v: i16) {
    b[o..o + 2].copy_from_slice(&v.to_le_bytes());
}

fn put_u32(b: &mut [u8], o: usize, v: u32) {
    b[o..o + 4].copy_from_slice(&v.to_le_bytes());
}
