// rasm-core/src/codegen/jit/mod.rs
//
//! # JIT-компилятор RASM
//!
//! Берёт готовый [`CompiledModule`] и выполняет его прямо в памяти.
//!
//! ## Раздельные страницы по секциям
//!
//! `.text` и `.data`/`.bss` получают **разные** страницы:
//! - `.text` — RX (исполняемый, доступный только на чтение после патчинга).
//! - `.data`/`.bss` — RW (доступный на запись, не исполняемый).
//!
//! Раньше данные директив (`.byte`/`.quad` и т.п.) физически лежали в том же
//! буфере, что код, и в итоге оказывались на странице без права записи —
//! любая попытка программы записать что-то в свою `.data` секцию во время
//! выполнения приводила бы к нарушению доступа.
//!
//! ## Как это устроено по платформам
//!
//! ### Linux / macOS
//! ```text
//! mmap(NULL, size, PROT_READ|PROT_WRITE, MAP_PRIVATE|MAP_ANON, -1, 0)
//!   -> страница rw-
//! memcpy(page, code, len)
//! mprotect(page, size, PROT_READ|PROT_EXEC)   ; только для .text
//!   -> страница rx-  (W^X)
//! munmap при дропе
//! ```
//!
//! ### Windows
//! ```text
//! VirtualAlloc(NULL, size, MEM_COMMIT|MEM_RESERVE, PAGE_READWRITE)
//!   -> страница rw-
//! memcpy(page, code, len)
//! VirtualProtect(page, size, PAGE_EXECUTE_READ, &old_protect)  ; только для .text
//!   -> страница rx-  (W^X)
//! VirtualFree при дропе
//! ```
//!
//! W^X (Write XOR Execute) — страница `.text` никогда не бывает одновременно
//! записываемой и исполняемой. `.data`/`.bss` остаются записываемыми и
//! никогда не получают права на исполнение.

pub mod symbol;

use crate::{
    codegen::{jit::symbol::SymbolTable, reloc::RelocKind},
    object::{CompiledModule, SectionId, SymbolVisibility},
};
use std::{collections::HashMap, io};

const EXTERNAL_THUNK_SIZE: usize = 12;

// =========================================================================
// Импорты Windows API
// =========================================================================

#[cfg(target_os = "windows")]
use windows_sys::Win32::System::Memory::{
    MEM_COMMIT, MEM_RELEASE, MEM_RESERVE, PAGE_EXECUTE_READ, PAGE_READWRITE, VirtualAlloc,
    VirtualFree, VirtualProtect,
};

// =========================================================================
// Ошибки
// =========================================================================

#[derive(Debug)]
pub enum JitError {
    /// Ошибка выделения исполняемой памяти.
    AllocFailed(io::Error),
    /// Точка входа не найдена в таблице меток.
    EntryNotFound(String),
    /// Внешний символ не удалось зарезолвить.
    UnresolvedSymbol(String),
}

impl std::fmt::Display for JitError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::AllocFailed(e) => write!(f, "JIT alloc failed: {e}"),
            Self::EntryNotFound(name) => write!(f, "entry point not found: '{name}'"),
            Self::UnresolvedSymbol(s) => write!(f, "unresolved external symbol: '{s}'"),
        }
    }
}

/// Результат выполнения JIT-кода.
pub type JitResult = Result<i64, JitError>;

// =========================================================================
// Одна выделенная страница памяти (используется и для .text, и для .data/.bss)
// =========================================================================

/// Одна выделенная через mmap/VirtualAlloc страница. Хранит **реальный**
/// выделенный размер (округлённый до страницы ОС), а не пересчитывает его
/// заново при освобождении — раньше `Drop` всегда вызывал
/// `munmap(page, page_size())` с размером **одной** страницы ОС, независимо
/// от того, сколько страниц было реально выделено через `round_to_page`,
/// что было утечкой памяти для любого буфера длиннее одной страницы.
struct JitPage {
    ptr: *mut u8,
    /// Реальный размер выделенного региона (кратен размеру страницы ОС).
    allocated_size: usize,
}

impl JitPage {
    /// Выделяет RW-страницу размера не менее `needed` байт (округляется
    /// вверх до кратного размеру страницы ОС) и копирует туда `data`,
    /// оставляя остаток обнулённым (актуально для `.bss`, где `data` может
    /// быть короче, чем зарезервированный размер).
    fn alloc_rw_with(needed: usize, data: &[u8]) -> Result<Self, JitError> {
        let allocated_size = round_to_page(needed.max(1));
        let ptr = alloc_rw(allocated_size).map_err(JitError::AllocFailed)?;

        unsafe {
            // mmap/VirtualAlloc уже отдают обнулённую память — обнулять
            // явно не нужно, копируем только реальные байты.
            if !data.is_empty() {
                std::ptr::copy_nonoverlapping(data.as_ptr(), ptr, data.len());
            }
        }

        Ok(Self {
            ptr,
            allocated_size,
        })
    }

    fn make_executable(&self) -> Result<(), JitError> {
        make_executable(self.ptr, self.allocated_size).map_err(JitError::AllocFailed)
    }
}

impl Drop for JitPage {
    fn drop(&mut self) {
        free_page(self.ptr, self.allocated_size);
    }
}

// =========================================================================
// JitExecutor
// =========================================================================

/// Владеет исполняемой страницей `.text` и, если они присутствуют в модуле,
/// записываемыми страницами `.data`/`.bss`. При дропе освобождает все
/// выделенные страницы через `munmap` (Unix) или `VirtualFree` (Windows).
pub struct JitExecutor {
    text_page: JitPage,
    /// Страницы остальных секций (`.data`, `.bss`, ...), которые получили
    /// память — держим их живыми ровно столько, сколько живёт исполняемый
    /// код, даже если сам JIT сейчас явно их не трогает после патчинга
    /// (программа обращается к ним по абсолютному адресу через `Abs64`).
    _data_pages: Vec<JitPage>,
    entry_offset: usize,
}

// SAFETY: JitExecutor владеет страницами эксклюзивно.
// После make_executable страница .text только читается + исполняется.
unsafe impl Send for JitExecutor {}

#[derive(Debug, Clone)]
struct ResolvedExternal {
    alias: String,
    addr: usize,
}

fn resolve_external(
    module: &CompiledModule,
    symbol_table: &mut SymbolTable,
    name: &str,
) -> Result<Option<ResolvedExternal>, JitError> {
    if let Some(import) = module.import(name) {
        let addr = symbol_table
            .resolve_import(&import.module, &import.name)
            .ok_or_else(|| {
                JitError::UnresolvedSymbol(format!("{}!{}", import.module, import.name))
            })?;
        return Ok(Some(ResolvedExternal {
            alias: import.alias.clone(),
            addr,
        }));
    }

    let Some(symbol) = module.symbol(name) else {
        return Ok(None);
    };

    if symbol.visibility != SymbolVisibility::External || symbol.section.is_some() {
        return Ok(None);
    }

    let addr = symbol_table
        .resolve(name)
        .ok_or_else(|| JitError::UnresolvedSymbol(name.to_string()))?;
    Ok(Some(ResolvedExternal {
        alias: name.to_string(),
        addr,
    }))
}

fn append_external_thunk(text: &mut Vec<u8>, addr: usize) -> usize {
    let offset = text.len();
    text.reserve(EXTERNAL_THUNK_SIZE);
    text.extend_from_slice(&[0x48, 0xB8]);
    text.extend_from_slice(&(addr as u64).to_le_bytes());
    text.extend_from_slice(&[0xFF, 0xE0]);
    debug_assert_eq!(text.len() - offset, EXTERNAL_THUNK_SIZE);
    offset
}

impl JitExecutor {
    /// Строит исполняемые/записываемые страницы из [`CompiledModule`],
    /// патчит `Abs64`-релокации (с учётом реальных адресов страниц,
    /// известных только после их выделения) и внешние символы.
    ///
    /// - `external_symbols` — список `(patch_offset_in_text, symbol_name)`
    ///   для резолвинга через `symbol_table` (offset — в байтах `.text`).
    pub fn new(
        module: &CompiledModule,
        symbol_table: &mut SymbolTable,
        external_symbols: &[(usize, String)],
    ) -> Result<Self, JitError> {
        let text_section = module
            .text_section()
            .ok_or_else(|| JitError::EntryNotFound("(no .text section)".to_string()))?;

        let entry_offset = module
            .entry
            .as_ref()
            .filter(|e| e.section == text_section.id)
            .map(|e| e.offset)
            .ok_or_else(|| JitError::EntryNotFound("_start/main".to_string()))?;

        let mut text_bytes = text_section.bytes.clone();
        let mut external_addrs: HashMap<String, usize> = HashMap::new();
        let mut thunk_offsets: HashMap<String, usize> = HashMap::new();

        for import in &module.imports {
            let resolved = resolve_external(module, symbol_table, &import.alias)?
                .ok_or_else(|| JitError::UnresolvedSymbol(import.alias.clone()))?;
            external_addrs.insert(resolved.alias.clone(), resolved.addr);
            let thunk_offset = append_external_thunk(&mut text_bytes, resolved.addr);
            thunk_offsets.insert(resolved.alias, thunk_offset);
        }

        for reloc in &module.relocations {
            if module.import(&reloc.target).is_some() || thunk_offsets.contains_key(&reloc.target) {
                continue;
            }

            if let Some(resolved) = resolve_external(module, symbol_table, &reloc.target)? {
                external_addrs.insert(resolved.alias.clone(), resolved.addr);
                let thunk_offset = append_external_thunk(&mut text_bytes, resolved.addr);
                thunk_offsets.insert(resolved.alias, thunk_offset);
            }
        }

        // 1. Выделяем RW-страницу под `.text` (патчим, потом переключим в RX).
        let text_page = JitPage::alloc_rw_with(text_bytes.len(), &text_bytes)?;

        // 2. Выделяем RW-страницы под все остальные непустые секции
        //    (`.data`, `.bss`, ...), в порядке `module.sections`. `.bss` не
        //    копирует байты сверх своего фактического содержимого — память
        //    под зарезервированный (но не записанный) хвост уже обнулена
        //    самим mmap/VirtualAlloc.
        let mut data_pages = Vec::new();
        // Параллельный список (SectionId, индекс в data_pages) — нужен для
        // патчинга релокаций, ссылающихся на эти секции.
        let mut section_pages: Vec<(crate::object::SectionId, usize)> = Vec::new();

        for section in &module.sections {
            if section.id == text_section.id || section.virtual_size == 0 {
                continue;
            }
            let page = JitPage::alloc_rw_with(section.virtual_size, &section.bytes)?;
            section_pages.push((section.id, data_pages.len()));
            data_pages.push(page);
        }

        // 3. Патчим `Abs64`-релокации: теперь, когда все страницы выделены,
        //    мы знаем их реальные адреса в памяти этого процесса.
        let page_addr_for = |section_id: SectionId| -> Option<*mut u8> {
            if section_id == text_section.id {
                Some(text_page.ptr)
            } else {
                section_pages
                    .iter()
                    .find(|(id, _)| *id == section_id)
                    .map(|(_, idx)| data_pages[*idx].ptr)
            }
        };

        for reloc in &module.relocations {
            let Some(patch_page_ptr) = page_addr_for(reloc.section) else {
                continue;
            };

            let target_addr = if let Some(&thunk_offset) = thunk_offsets.get(&reloc.target) {
                unsafe { text_page.ptr.add(thunk_offset) as usize }
            } else if let Some(symbol) = module.symbol(&reloc.target) {
                if let Some(target_section) = symbol.section {
                    let Some(target_page_ptr) = page_addr_for(target_section) else {
                        continue;
                    };
                    unsafe { target_page_ptr.add(symbol.offset) as usize }
                } else if let Some(&addr) = external_addrs.get(&reloc.target) {
                    addr
                } else {
                    continue;
                }
            } else if let Some(&addr) = external_addrs.get(&reloc.target) {
                addr
            } else {
                continue;
            };

            match reloc.kind {
                RelocKind::Abs64 => unsafe {
                    let patch_ptr = patch_page_ptr.add(reloc.offset) as *mut u64;
                    std::ptr::write_unaligned(patch_ptr, target_addr as u64);
                },
                RelocKind::Rel32 => unsafe {
                    let patch_addr = patch_page_ptr.add(reloc.offset) as usize;
                    let rel = target_addr as i128 - (patch_addr + 4) as i128;
                    if rel < i32::MIN as i128 || rel > i32::MAX as i128 {
                        return Err(JitError::UnresolvedSymbol(format!(
                            "{}: rel32 target is out of range",
                            reloc.target
                        )));
                    }
                    let patch_ptr = patch_page_ptr.add(reloc.offset) as *mut i32;
                    std::ptr::write_unaligned(patch_ptr, rel as i32);
                },
            }
        }

        // 4. Патчим внешние символы (всегда в `.text`, как и раньше).
        for (patch_offset, sym_name) in external_symbols {
            let addr = symbol_table
                .resolve(sym_name)
                .ok_or_else(|| JitError::UnresolvedSymbol(sym_name.clone()))?;
            unsafe {
                let ptr = text_page.ptr.add(*patch_offset) as *mut u64;
                std::ptr::write_unaligned(ptr, addr as u64);
            }
        }

        // 5. Переключаем `.text` в rx- (W^X). `.data`/`.bss` остаются rw-.
        text_page.make_executable()?;

        Ok(Self {
            text_page,
            _data_pages: data_pages,
            entry_offset,
        })
    }

    /// Выполняет JIT-код как `fn() -> i64`.
    ///
    /// # Safety
    /// Буфер должен содержать валидный x86_64 машинный код.
    /// Вызывающая сторона отвечает за корректность кода.
    pub unsafe fn execute(&self) -> i64 {
        unsafe {
            let fn_ptr: extern "C" fn() -> i64 =
                std::mem::transmute(self.text_page.ptr.add(self.entry_offset));
            fn_ptr()
        }
    }

    /// Выполняет JIT-код как `fn(i64) -> i64`.
    /// # Safety
    ///
    /// Вызывающий код должен гарантировать, что:
    /// 1. Память по адресу страницы `.text` содержит валидный и полностью
    ///    скомпилированный исполняемый машинный код (C-ABI).
    /// 2. `self.entry_offset` указывает строго на начало функции.
    /// 3. Страница памяти имеет права на выполнение (rx-/rwx).
    /// 4. Переданный аргумент `arg` и возвращаемое значение не нарушают
    ///    инварианты JIT-кода.
    pub unsafe fn execute_with_arg(&self, arg: i64) -> i64 {
        unsafe {
            let fn_ptr: extern "C" fn(i64) -> i64 =
                std::mem::transmute(self.text_page.ptr.add(self.entry_offset));
            fn_ptr(arg)
        }
    }
}

// =========================================================================
// Платформо-зависимые примитивы
// =========================================================================

/// Округляет `needed` вверх до кратного размеру страницы ОС.
fn round_to_page(needed: usize) -> usize {
    let page = page_size();
    (needed + page - 1) & !(page - 1)
}

// ── Размер страницы ───────────────────────────────────────────────────

#[cfg(unix)]
fn page_size() -> usize {
    unsafe { libc::sysconf(libc::_SC_PAGESIZE) as usize }
}

#[cfg(target_os = "windows")]
fn page_size() -> usize {
    use windows_sys::Win32::System::SystemInformation::{GetSystemInfo, SYSTEM_INFO};
    unsafe {
        let mut info: SYSTEM_INFO = std::mem::zeroed();
        GetSystemInfo(&mut info);
        info.dwPageSize as usize
    }
}

#[cfg(not(any(unix, target_os = "windows")))]
fn page_size() -> usize {
    4096
}

// ── Выделение rw- страницы ────────────────────────────────────────────

#[cfg(unix)]
fn alloc_rw(size: usize) -> io::Result<*mut u8> {
    let ptr = unsafe {
        libc::mmap(
            std::ptr::null_mut(),
            size,
            libc::PROT_READ | libc::PROT_WRITE,
            libc::MAP_PRIVATE | libc::MAP_ANONYMOUS,
            -1,
            0,
        )
    };
    if ptr == libc::MAP_FAILED {
        Err(io::Error::last_os_error())
    } else {
        Ok(ptr as *mut u8)
    }
}

#[cfg(target_os = "windows")]
fn alloc_rw(size: usize) -> io::Result<*mut u8> {
    let ptr = unsafe {
        VirtualAlloc(
            std::ptr::null(), // адрес: ОС выберет сама
            size,
            MEM_COMMIT | MEM_RESERVE, // сразу коммитим физическую память
            PAGE_READWRITE,           // rw-, без exec
        )
    };
    if ptr.is_null() {
        Err(io::Error::last_os_error())
    } else {
        Ok(ptr as *mut u8)
    }
}

#[cfg(not(any(unix, target_os = "windows")))]
fn alloc_rw(_size: usize) -> io::Result<*mut u8> {
    Err(io::Error::new(
        io::ErrorKind::Unsupported,
        "JIT is not supported on this platform",
    ))
}

// ── Переключение в rx- (W^X) ──────────────────────────────────────────

#[cfg(unix)]
fn make_executable(page: *mut u8, size: usize) -> io::Result<()> {
    let ret = unsafe {
        libc::mprotect(
            page as *mut libc::c_void,
            size,
            libc::PROT_READ | libc::PROT_EXEC,
        )
    };
    if ret != 0 {
        Err(io::Error::last_os_error())
    } else {
        Ok(())
    }
}

#[cfg(target_os = "windows")]
fn make_executable(page: *mut u8, size: usize) -> io::Result<()> {
    let mut old_protect: u32 = 0;
    let ret = unsafe {
        VirtualProtect(
            page as *const _,
            size,
            PAGE_EXECUTE_READ, // rx-, без write
            &mut old_protect,
        )
    };
    if ret == 0 {
        Err(io::Error::last_os_error())
    } else {
        Ok(())
    }
}

#[cfg(not(any(unix, target_os = "windows")))]
fn make_executable(_page: *mut u8, _size: usize) -> io::Result<()> {
    Ok(())
}

// ── Освобождение страницы ─────────────────────────────────────────────
//
// Принимает реальный выделенный размер (см. `JitPage::allocated_size`),
// а не пересчитывает размер страницы ОС заново — раньше освобождался
// ровно один `page_size()` независимо от того, сколько страниц было
// реально выделено для буфера длиннее одной страницы (утечка памяти).

#[cfg(unix)]
fn free_page(ptr: *mut u8, allocated_size: usize) {
    unsafe {
        libc::munmap(ptr as *mut libc::c_void, allocated_size);
    }
}

#[cfg(target_os = "windows")]
fn free_page(ptr: *mut u8, _allocated_size: usize) {
    unsafe {
        // MEM_RELEASE требует size = 0 при освобождении всего региона,
        // выделенного одним вызовом VirtualAlloc.
        VirtualFree(ptr as *mut _, 0, MEM_RELEASE);
    }
}

#[cfg(not(any(unix, target_os = "windows")))]
fn free_page(_ptr: *mut u8, _allocated_size: usize) {}
