// rasm-core/src/codegen/mod.rs
//
//! # Кодогенератор RASM
//!
//! Принимает [`AST`] и превращает его в плоский байтовый буфер — машинный код x86_64.
//!
//! ## Алгоритм работы
//!
//! ```text
//! AST
//!  │
//!  ├─ Pass 1: обход Statement-ов, кодирование инструкций
//!  │    ├─ встретили Label                            -> записываем label_name -> текущий offset в label_offsets
//!  │    ├─ встретили Instruction                      -> encode_instruction()  -> байты в Emitter
//!  │    │    └─ если операнд Label и адрес неизвестен -> 4 нулевых байта + запись в RelocTable
//!  │    └─ встретили Directive                        -> обрабатываем (data секция, align и т.д.)
//!  │
//!  └─ Pass 2: патчинг релокаций
//!       └─ для каждой Relocation в RelocTable:
//!            ├─ ищем label в label_offsets
//!            ├─ вычисляем rel32 = target - (patch_offset + 4)
//!            └─ emit.patch_i32(patch_offset, rel32)
//! ```
//!
//! ## Секции
//!
//! Пока кодогенератор работает в режиме "flat binary": все инструкции идут подряд,
//! данные из `.byte`/`.quad` и т.д. записываются прямо в тот же буфер в том порядке,
//! в котором встретились. Это подходит для JIT.
//!
//! AOT (ELF/PE) потребует разделения на `.text` и `.data` секции — это в `aot/`.
pub mod aot;
pub mod emit;
pub mod encode;
pub mod err;
pub mod jit;
pub mod reloc;

use std::collections::HashMap;

use crate::{
    codegen::{
        emit::Emitter,
        encode::encode_instruction,
        err::{CodegenError, CodegenErrorKind},
        reloc::{RelocKind, RelocTable},
    },
    isa::directive::{Directive, DirectiveKind, DirectiveOperand},
    object::{
        CompiledModule, Import, ObjectRelocation, Section, SectionId, SectionKind, Symbol,
        SymbolVisibility,
    },
    parser::ast::{AST, Statement},
    str_pool::{StrId, StrPool},
};

/// # Результат кодогенерации
///
/// Содержит готовый байтовый буфер и карту меток (нужна AOT и JIT).
#[derive(Debug)]
pub struct CodegenOutput {
    /// Плоский байтовый буфер с машинным кодом.
    pub code: Vec<u8>,

    /// Карта `имя_метки -> байтовый_офсет` в `code`.
    /// Нужна JIT-компилятору, чтобы найти точку входа (`_start`, `main`).
    pub label_offsets: HashMap<String, usize>,

    /// Структурированное представление результата сборки.
    pub module: CompiledModule,
}

#[derive(Debug, Clone)]
struct FlatRange {
    flat_start: usize,
    section_start: usize,
    len: usize,
}

#[derive(Debug, Clone)]
struct SectionBuilder {
    id: SectionId,
    name: &'static str,
    kind: SectionKind,
    bytes: Vec<u8>,
    align: usize,
    ranges: Vec<FlatRange>,
}

impl SectionBuilder {
    fn new(id: SectionId, name: &'static str, kind: SectionKind, align: usize) -> Self {
        Self {
            id,
            name,
            kind,
            bytes: Vec::new(),
            align,
            ranges: Vec::new(),
        }
    }

    fn offset(&self) -> usize {
        self.bytes.len()
    }

    fn append_from_flat(&mut self, flat_start: usize, bytes: &[u8]) {
        if bytes.is_empty() {
            return;
        }

        let section_start = self.bytes.len();
        self.bytes.extend_from_slice(bytes);
        self.ranges.push(FlatRange {
            flat_start,
            section_start,
            len: bytes.len(),
        });
    }

    fn sync_from_flat(&mut self, flat: &[u8]) {
        for range in &self.ranges {
            let flat_end = range.flat_start + range.len;
            let section_end = range.section_start + range.len;
            self.bytes[range.section_start..section_end]
                .copy_from_slice(&flat[range.flat_start..flat_end]);
        }
    }

    fn into_section(self) -> Section {
        Section::new(self.id, self.name, self.kind, self.bytes, self.align)
    }
}

/// # Кодогенератор
///
/// Создаётся один раз на всю компиляцию. Хранит внутри себя [`Emitter`] и [`RelocTable`].
pub struct CodeGen<'a> {
    emit: Emitter,
    relocs: RelocTable,
    /// Карта StrId -> байтовый офсет метки. Заполняется по ходу первого прохода.
    label_offsets: HashMap<StrId, usize>,
    label_locations: HashMap<StrId, (SectionId, usize)>,
    symbol_visibility: HashMap<StrId, SymbolVisibility>,
    imports: HashMap<StrId, Import>,
    sections: Vec<SectionBuilder>,
    current_section: SectionId,
    str_pool: &'a StrPool,
    /// Накопленные ошибки. Кодогенератор не останавливается на первой ошибке.
    pub errors: Vec<CodegenError>,
}

impl<'a> CodeGen<'a> {
    pub fn new(str_pool: &'a StrPool) -> Self {
        Self {
            emit: Emitter::with_capacity(1024),
            relocs: RelocTable::new(),
            label_offsets: HashMap::new(),
            label_locations: HashMap::new(),
            symbol_visibility: HashMap::new(),
            imports: HashMap::new(),
            sections: vec![SectionBuilder::new(
                SectionId(0),
                ".text",
                SectionKind::Text,
                16,
            )],
            current_section: SectionId(0),
            str_pool,
            errors: Vec::new(),
        }
    }

    /// Запускает кодогенерацию для всего AST и возвращает [`CodegenOutput`].
    ///
    /// Если во время кодирования встретились ошибки, они будут в `self.errors`.
    /// `code` при этом всё равно возвращается — с теми байтами, что удалось сгенерировать.
    pub fn compile(&mut self, ast: &AST) -> CodegenOutput {
        // --- Pass 1: кодируем инструкции, параллельно собираем label_offsets ---
        for stmt in &ast.program {
            match stmt {
                Statement::Label { span: _, name: id } => {
                    // Метка -> запоминаем текущий офсет
                    self.label_offsets.insert(*id, self.emit.offset());
                    self.label_locations
                        .insert(*id, (self.current_section, self.current_section_offset()));
                }

                Statement::Instruction {
                    span: _,
                    instruction: instr,
                } => {
                    let before = self.emit.offset();
                    if let Err(e) = encode_instruction(
                        instr,
                        &mut self.emit,
                        &mut self.relocs,
                        &self.label_offsets,
                    ) {
                        self.errors.push(e);
                    }
                    self.record_emitted_range(before);
                }

                Statement::Directive {
                    span: _,
                    directive: dir,
                } => {
                    self.compile_directive(dir);
                }
            }
        }

        // --- Pass 2: патчим релокации (forward references) ---
        self.apply_relocations();
        self.sync_sections_from_flat();

        // --- Строим финальный output ---
        let code = self.emit.buf.clone();
        let label_offsets = self
            .label_offsets
            .iter()
            .map(|(id, &off)| (self.str_pool.get(*id).to_string(), off))
            .collect();
        let module = self.build_module();

        CodegenOutput {
            code,
            label_offsets,
            module,
        }
    }

    // =========================================================================
    // Директивы
    // =========================================================================

    fn compile_directive(&mut self, dir: &Directive) {
        match dir.kind {
            DirectiveKind::Text => self.switch_section(".text", SectionKind::Text, 16),
            DirectiveKind::Data => self.switch_section(".data", SectionKind::Data, 8),
            DirectiveKind::Bss => self.switch_section(".bss", SectionKind::Bss, 8),

            // .global / .extern — информация для линкера, не генерирует байты
            DirectiveKind::Global => {
                self.record_symbol_visibility(dir, SymbolVisibility::Global);
            }
            DirectiveKind::Extern => {
                self.record_symbol_visibility(dir, SymbolVisibility::External);
            }
            DirectiveKind::Import => {
                self.record_import(dir);
            }

            // .byte — 1 байт на каждый операнд
            DirectiveKind::Byte => {
                let before = self.emit.offset();
                for op in &dir.operands {
                    if let DirectiveOperand::Number(n) = op {
                        self.emit.emit_u8(*n as u8);
                    }
                }
                self.record_emitted_range(before);
            }

            // .word — 2 байта (little-endian)
            DirectiveKind::Word => {
                let before = self.emit.offset();
                for op in &dir.operands {
                    if let DirectiveOperand::Number(n) = op {
                        self.emit.emit_u16(*n as u16);
                    }
                }
                self.record_emitted_range(before);
            }

            // .dword — 4 байта
            DirectiveKind::Dword => {
                let before = self.emit.offset();
                for op in &dir.operands {
                    if let DirectiveOperand::Number(n) = op {
                        self.emit.emit_u32(*n as u32);
                    }
                }
                self.record_emitted_range(before);
            }

            // .quad — 8 байт; если операнд — символ, добавляем Abs64 релокацию
            DirectiveKind::Qword => {
                let before = self.emit.offset();
                for op in &dir.operands {
                    match op {
                        DirectiveOperand::Number(n) => {
                            self.emit.emit_i64(*n);
                        }
                        DirectiveOperand::Symbol(id) => {
                            let patch_offset = self.emit.offset();
                            self.emit.emit_u64(0); // заглушка
                            self.relocs.add(patch_offset, *id, RelocKind::Abs64);
                        }
                        _ => {}
                    }
                }
                self.record_emitted_range(before);
            }

            // .ascii — записываем байты строки как есть
            DirectiveKind::Ascii => {
                let before = self.emit.offset();
                for op in &dir.operands {
                    if let DirectiveOperand::StringLiteral(id) = op {
                        let s = self.str_pool.get(*id);
                        self.emit.emit_bytes(s.as_bytes());
                    }
                }
                self.record_emitted_range(before);
            }

            // .align N — выравниваем текущий офсет до кратного N
            // В .text заполняем NOP (0x90), в других секциях — нулями.
            // Здесь используем NOP как универсальный вариант.
            DirectiveKind::Align => {
                if let Some(DirectiveOperand::Number(n)) = dir.operands.first() {
                    let align = *n as usize;
                    if align > 1 && align.is_power_of_two() {
                        let before = self.emit.offset();
                        let current = self.current_section_offset();
                        let remainder = current % align;
                        if remainder != 0 {
                            let padding = align - remainder;
                            let fill = if self.current_section_kind() == SectionKind::Text {
                                0x90
                            } else {
                                0x00
                            };
                            for _ in 0..padding {
                                self.emit.emit_u8(fill);
                            }
                            self.record_emitted_range(before);
                        }
                    }
                }
            }
        }
    }

    fn switch_section(&mut self, name: &'static str, kind: SectionKind, align: usize) {
        if let Some(section) = self.sections.iter().find(|section| section.name == name) {
            self.current_section = section.id;
            return;
        }

        let id = SectionId(self.sections.len());
        self.sections
            .push(SectionBuilder::new(id, name, kind, align));
        self.current_section = id;
    }

    fn current_section(&self) -> &SectionBuilder {
        &self.sections[self.current_section.0]
    }

    fn current_section_mut(&mut self) -> &mut SectionBuilder {
        &mut self.sections[self.current_section.0]
    }

    fn current_section_offset(&self) -> usize {
        self.current_section().offset()
    }

    fn current_section_kind(&self) -> SectionKind {
        self.current_section().kind
    }

    fn record_emitted_range(&mut self, flat_start: usize) {
        let flat_end = self.emit.offset();
        if flat_end == flat_start {
            return;
        }

        let bytes = self.emit.buf[flat_start..flat_end].to_vec();
        self.current_section_mut()
            .append_from_flat(flat_start, &bytes);
    }

    fn sync_sections_from_flat(&mut self) {
        for section in &mut self.sections {
            section.sync_from_flat(&self.emit.buf);
        }
    }

    fn record_symbol_visibility(&mut self, dir: &Directive, visibility: SymbolVisibility) {
        for op in &dir.operands {
            if let DirectiveOperand::Symbol(id) = op {
                self.symbol_visibility.insert(*id, visibility);
            }
        }
    }

    fn record_import(&mut self, dir: &Directive) {
        let [
            DirectiveOperand::StringLiteral(module_id),
            DirectiveOperand::Symbol(name_id),
        ] = dir.operands.as_slice()
        else {
            return;
        };

        let module = self.str_pool.get(*module_id).to_string();
        let name = self.str_pool.get(*name_id).to_string();
        self.symbol_visibility
            .insert(*name_id, SymbolVisibility::External);
        self.imports
            .insert(*name_id, Import::new(module, name.clone(), name));
    }

    fn build_module(&self) -> CompiledModule {
        let sections = self
            .sections
            .iter()
            .cloned()
            .map(SectionBuilder::into_section)
            .collect();

        let mut symbols = self
            .label_locations
            .iter()
            .map(|(id, &(section, offset))| {
                let visibility = self
                    .symbol_visibility
                    .get(id)
                    .copied()
                    .filter(|visibility| *visibility != SymbolVisibility::External)
                    .unwrap_or(SymbolVisibility::Local);
                Symbol::defined(
                    self.str_pool.get(*id).to_string(),
                    section,
                    offset,
                    visibility,
                )
            })
            .collect::<Vec<_>>();

        for (&id, &visibility) in &self.symbol_visibility {
            if visibility == SymbolVisibility::External && !self.label_locations.contains_key(&id) {
                symbols.push(Symbol::external(self.str_pool.get(id).to_string()));
            }
        }

        let relocations = self
            .relocs
            .entries
            .iter()
            .map(|reloc| {
                let (section, offset) = self
                    .flat_to_section_offset(reloc.patch_offset)
                    .unwrap_or((SectionId(0), reloc.patch_offset));
                ObjectRelocation {
                    section,
                    offset,
                    kind: reloc.kind,
                    target: self.str_pool.get(reloc.target_label).to_string(),
                }
            })
            .collect();

        let imports = self.imports.values().cloned().collect();

        CompiledModule::with_imports(sections, symbols, relocations, imports)
    }

    fn flat_to_section_offset(&self, flat_offset: usize) -> Option<(SectionId, usize)> {
        for section in &self.sections {
            for range in &section.ranges {
                if (range.flat_start..range.flat_start + range.len).contains(&flat_offset) {
                    return Some((
                        section.id,
                        range.section_start + (flat_offset - range.flat_start),
                    ));
                }
            }
        }

        None
    }

    // =========================================================================
    // Патчинг релокаций (Pass 2)
    // =========================================================================

    fn apply_relocations(&mut self) {
        for reloc in &self.relocs.entries {
            match self.label_offsets.get(&reloc.target_label) {
                None => {
                    if self.symbol_visibility.get(&reloc.target_label)
                        == Some(&SymbolVisibility::External)
                    {
                        continue;
                    }

                    // Метка не найдена — ошибка
                    let name = self.str_pool.get(reloc.target_label).to_string();
                    self.errors
                        .push(CodegenError::new(CodegenErrorKind::UndefinedLabel(name)));
                }

                Some(&target_offset) => match reloc.kind {
                    RelocKind::Rel32 => {
                        // rel32 = target - (patch_offset + 4)
                        let rel = (target_offset as i64) - (reloc.patch_offset as i64 + 4);

                        if rel > i32::MAX as i64 || rel < i32::MIN as i64 {
                            let name = self.str_pool.get(reloc.target_label).to_string();
                            self.errors.push(CodegenError::new(
                                CodegenErrorKind::RelocationOverflow {
                                    label: name,
                                    offset: rel,
                                },
                            ));
                        } else {
                            self.emit.patch_i32(reloc.patch_offset, rel as i32);
                        }
                    }

                    // `Abs64` НЕ патчится здесь. Финальный адрес зависит от
                    // базы загрузки, которую кодогенератор не знает (JIT
                    // мапит код по произвольному адресу, AOT грузит по
                    // фиксированной базе вида 0x400000) — то, что раньше
                    // здесь предполагалось как "base = 0, JIT досчитает",
                    // было неверно для AOT: ELF/PE писали этот же offset
                    // в файл без поправки на базу образа, давая неверный
                    // абсолютный адрес в `.quad label`.
                    //
                    // Вместо патчинга здесь заглушка (8 нулевых байт,
                    // записанная в `compile_directive` при `.quad label`)
                    // остаётся как есть, а сама релокация уходит в
                    // `module.relocations` (см. `build_module`) — каждый
                    // потребитель (`JitExecutor`, `ElfWriter`, `PeWriter`)
                    // патчит её сам, зная свою базу загрузки.
                    RelocKind::Abs64 => {}
                },
            }
        }
    }
}
