//! Резолвинг внешних символов для JIT.
//!
//! `.extern foo` ищется в уже загруженных символах процесса, а
//! `.import module foo` сначала загружает указанный модуль и потом ищет
//! функцию внутри него.

use std::{collections::HashMap, ffi::CString};

#[cfg(target_os = "windows")]
use windows_sys::Win32::System::LibraryLoader::{GetProcAddress, LoadLibraryA};

#[cfg(unix)]
type LibraryRawHandle = *mut libc::c_void;

#[cfg(target_os = "windows")]
type LibraryRawHandle = windows_sys::Win32::Foundation::HMODULE;

#[cfg(not(any(unix, target_os = "windows")))]
type LibraryRawHandle = ();

#[derive(Debug)]
struct LibraryHandle {
    handle: LibraryRawHandle,
}

/// Таблица внешних символов, доступных JIT-коду.
#[derive(Debug, Default)]
pub struct SymbolTable {
    resolved: HashMap<String, usize>,
    libraries: HashMap<String, LibraryHandle>,
}

impl SymbolTable {
    pub fn new() -> Self {
        Self {
            resolved: HashMap::new(),
            libraries: HashMap::new(),
        }
    }

    /// Ищет символ в текущем процессе или среди вручную зарегистрированных адресов.
    pub fn resolve(&mut self, name: &str) -> Option<usize> {
        if let Some(&addr) = self.resolved.get(name) {
            return Some(addr);
        }

        let addr = self.dlsym_lookup(name)?;
        self.resolved.insert(name.to_string(), addr);
        Some(addr)
    }

    /// Ищет символ в конкретной динамической библиотеке.
    pub fn resolve_import(&mut self, module: &str, name: &str) -> Option<usize> {
        let key = format!("{module}!{name}");
        if let Some(&addr) = self.resolved.get(&key) {
            return Some(addr);
        }

        let handle = self.load_library(module)?;
        let addr = lookup_in_library(handle, name)?;
        self.resolved.insert(key, addr);
        Some(addr)
    }

    /// Регистрирует символ вручную, например Rust-функцию для тестов или embedding API.
    ///
    /// ```ignore
    /// extern "C" fn my_print(n: i64) { println!("{n}"); }
    /// symbol_table.register("my_print", my_print as usize);
    /// ```
    pub fn register(&mut self, name: &str, addr: usize) {
        self.resolved.insert(name.to_string(), addr);
    }

    fn load_library(&mut self, module: &str) -> Option<LibraryRawHandle> {
        let key = module.to_ascii_lowercase();
        if let Some(handle) = self.libraries.get(&key) {
            return Some(handle.handle);
        }

        let handle = open_library(module)?;
        self.libraries.insert(key.clone(), LibraryHandle { handle });
        Some(self.libraries.get(&key)?.handle)
    }

    /// Выполняет поиск через `dlsym(RTLD_DEFAULT, name)`.
    #[cfg(unix)]
    fn dlsym_lookup(&self, name: &str) -> Option<usize> {
        unsafe {
            let c_name = CString::new(name).ok()?;
            let ptr = libc::dlsym(libc::RTLD_DEFAULT, c_name.as_ptr());
            if ptr.is_null() {
                None
            } else {
                Some(ptr as usize)
            }
        }
    }

    #[cfg(target_os = "windows")]
    fn dlsym_lookup(&self, name: &str) -> Option<usize> {
        self.resolved.get(name).copied()
    }

    #[cfg(not(any(unix, target_os = "windows")))]
    const fn dlsym_lookup(&self, _name: &str) -> Option<usize> {
        None
    }
}

#[cfg(unix)]
fn open_library(module: &str) -> Option<LibraryRawHandle> {
    unsafe {
        let c_module = CString::new(module).ok()?;
        let handle = libc::dlopen(c_module.as_ptr(), libc::RTLD_NOW);
        if handle.is_null() { None } else { Some(handle) }
    }
}

#[cfg(target_os = "windows")]
fn open_library(module: &str) -> Option<LibraryRawHandle> {
    unsafe {
        let c_module = CString::new(module).ok()?;
        let handle = LoadLibraryA(c_module.as_ptr() as *const u8);
        if handle.is_null() { None } else { Some(handle) }
    }
}

#[cfg(not(any(unix, target_os = "windows")))]
fn open_library(_module: &str) -> Option<LibraryRawHandle> {
    None
}

#[cfg(unix)]
fn lookup_in_library(handle: LibraryRawHandle, name: &str) -> Option<usize> {
    unsafe {
        let c_name = CString::new(name).ok()?;
        let ptr = libc::dlsym(handle, c_name.as_ptr());
        if ptr.is_null() {
            None
        } else {
            Some(ptr as usize)
        }
    }
}

#[cfg(target_os = "windows")]
fn lookup_in_library(handle: LibraryRawHandle, name: &str) -> Option<usize> {
    unsafe {
        let c_name = CString::new(name).ok()?;
        GetProcAddress(handle, c_name.as_ptr() as *const u8).map(|proc| proc as usize)
    }
}

#[cfg(not(any(unix, target_os = "windows")))]
fn lookup_in_library(_handle: LibraryRawHandle, _name: &str) -> Option<usize> {
    None
}
