// 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},
    parser::ast::{AST, Statement},
    str_pool::StrPool,
};

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

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

/// # Кодогенератор
///
/// Создаётся один раз на всю компиляцию. Хранит внутри себя [`Emitter`] и [`RelocTable`].
pub struct CodeGen<'a> {
    emit: Emitter,
    relocs: RelocTable,
    /// Карта StrId -> байтовый офсет метки. Заполняется по ходу первого прохода.
    label_offsets: HashMap<crate::str_pool::StrId, usize>,
    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(),
            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(id) => {
                    // Метка -> запоминаем текущий офсет
                    self.label_offsets.insert(*id, self.emit.offset());
                }

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

                Statement::Directive(dir) => {
                    self.compile_directive(dir);
                }
            }
        }

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

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

        CodegenOutput {
            code: self.emit.buf.clone(),
            label_offsets,
        }
    }

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

    fn compile_directive(&mut self, dir: &Directive) {
        match dir.kind {
            // Секции — в flat binary режиме просто игнорируем переключение
            DirectiveKind::Text | DirectiveKind::Data | DirectiveKind::Bss => {}

            // .global / .extern — информация для линкера, не генерирует байты
            DirectiveKind::Global | DirectiveKind::Extern => {}

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

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

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

            // .quad — 8 байт; если операнд — символ, добавляем Abs64 релокацию
            DirectiveKind::Qword => {
                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);
                        }
                        _ => {}
                    }
                }
            }

            // .ascii — записываем байты строки как есть
            DirectiveKind::Ascii => {
                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());
                        // Примечание: .asciz добавил бы нулевой байт в конце.
                        // Если нужно — проверяй исходный синтаксис (.asciz vs .ascii).
                    }
                }
            }

            // .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 current = self.emit.offset();
                        let remainder = current % align;
                        if remainder != 0 {
                            let padding = align - remainder;
                            for _ in 0..padding {
                                self.emit.emit_u8(0x90); // NOP
                            }
                        }
                    }
                }
            }
        }
    }

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

    fn apply_relocations(&mut self) {
        for reloc in &self.relocs.entries {
            match self.label_offsets.get(&reloc.target_label) {
                None => {
                    // Метка не найдена — ошибка
                    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);
                        }
                    }

                    RelocKind::Abs64 => {
                        // Для JIT: абсолютный адрес = base + target_offset.
                        // В flat binary base = 0, пишем просто offset.
                        // JIT-рантайм потом прибавит реальный базовый адрес при маппинге.
                        self.emit
                            .patch_u32(reloc.patch_offset, target_offset as u32);
                    }
                },
            }
        }
    }
}
