// rasm-core/src/codegen/jit/mod.rs
//
//! # JIT-компилятор RASM
//!
//! Берёт готовый байтовый буфер из [`CodeGen`] и выполняет его прямо в памяти.
//!
//! ## Как это устроено по платформам
//!
//! ### 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)
//!   -> страница 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)
//!   -> страница rx-  (W^X)
//! VirtualFree при дропе
//! ```
//!
//! W^X (Write XOR Execute) — страница никогда не бывает одновременно
//! записываемой и исполняемой. Это базовая защита от эксплойтов.

pub mod symbol;

use crate::codegen::jit::symbol::SymbolTable;
use std::io;

// =========================================================================
// Импорты 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>;

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

/// Владеет исполняемой страницей памяти.
/// При дропе освобождает её через `munmap` (Unix) или `VirtualFree` (Windows).
pub struct JitExecutor {
    page: *mut u8,
    entry_offset: usize,
}

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

impl JitExecutor {
    /// Создаёт исполняемую страницу, копирует `code` и патчит внешние символы.
    ///
    /// - `entry_offset` — байтовый офсет точки входа внутри `code`
    /// - `external_symbols` — список `(patch_offset, symbol_name)` для резолвинга
    pub fn new(
        code: &[u8],
        entry_offset: usize,
        symbol_table: &mut SymbolTable,
        external_symbols: &[(usize, String)],
    ) -> Result<Self, JitError> {
        let page_size = Self::round_to_page(code.len());

        // 1. Выделяем страницу rw-
        let page = Self::alloc_rw(page_size).map_err(JitError::AllocFailed)?;

        // 2. Копируем байты
        unsafe {
            std::ptr::copy_nonoverlapping(code.as_ptr(), page, code.len());
        }

        // 3. Патчим внешние символы
        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 = page.add(*patch_offset) as *mut u64;
                *ptr = addr as u64;
            }
        }

        // 4. Переключаем в rx- (W^X)
        Self::make_executable(page, page_size).map_err(JitError::AllocFailed)?;

        Ok(Self { page, 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.page.add(self.entry_offset));
            fn_ptr()
        }
    }

    /// Выполняет JIT-код как `fn(i64) -> i64`.
    /// # Safety
    ///
    /// Вызывающий код должен гарантировать, что:
    /// 1. Память по адресу `self.page` содержит валидный и полностью скомпилированный
    ///    исполняемый машинный код (соответствующий соглашению о вызовах C-ABI).
    /// 2. `self.entry_offset` указывает строго на начало функции внутри этой страницы.
    /// 3. Страница памяти имеет права на выполнение (вызван `mprotect` с 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.page.add(self.entry_offset));
            fn_ptr(arg)
        }
    }

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

    /// Округляет `needed` вверх до кратного размеру страницы ОС.
    fn round_to_page(needed: usize) -> usize {
        let page = Self::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(())
    }
}

// =========================================================================
// Освобождение памяти
// =========================================================================

impl Drop for JitExecutor {
    fn drop(&mut self) {
        #[cfg(unix)]
        unsafe {
            libc::munmap(self.page as *mut libc::c_void, JitExecutor::page_size());
        }

        #[cfg(target_os = "windows")]
        unsafe {
            // MEM_RELEASE требует size = 0 при освобождении всего региона
            VirtualFree(self.page as *mut _, 0, MEM_RELEASE);
        }
    }
}
