//! Core hooking logic and types.
//!
//! Provides the [HookRegistry] for managing active hooks, and defines
//! the variants and instances of hooks.

use crate::error::PrehookError;
use goblin::elf::Elf;
use ilhook::x64::HookType as IlHookType;
use ilhook::x64::{CallbackOption, HookFlags, HookPoint, Hooker};
pub use ilhook::x64::{HookType, Registers};
pub use ilhook::x64::{JmpBackRoutine, JmpToAddrRoutine, JmpToRetRoutine, RetnRoutine};
use std::collections::HashMap;
use std::fs;
use std::sync::{Mutex, OnceLock};
use tracing::{debug, error, warn};

thread_local! {
    static HOOK_GUARD: std::cell::RefCell<Vec<&'static str>> = const { std::cell::RefCell::new(Vec::new()) };
}

pub struct HookGuard;

impl HookGuard {
    pub fn try_enter(name: &'static str) -> bool {
        HOOK_GUARD.with(|guard| {
            let mut guard = guard.borrow_mut();
            if guard.contains(&name) {
                false
            } else {
                guard.push(name);
                true
            }
        })
    }

    pub fn exit(name: &'static str) {
        HOOK_GUARD.with(|guard| {
            let mut guard = guard.borrow_mut();
            if let Some(pos) = guard.iter().rposition(|&n| n == name) {
                guard.remove(pos);
            }
        })
    }

    pub fn is_active(name: &'static str) -> bool {
        HOOK_GUARD.with(|guard| guard.borrow().contains(&name))
    }
}

/// Target specification for hook source.
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Target {
    /// Raw offset from binary base
    Raw(usize),
    /// Symbol name only
    Symbol(String),
    /// Symbol with offset: "func+0x10" or "func-0x5"
    SymbolOffset(String, isize),
    /// Offset relative to hook point: "+0x20" (for jmp_to_addr dest only)
    Relative(isize),
}

impl Target {
    pub fn resolve(&self, base: usize, hook_point: usize) -> Result<usize, PrehookError> {
        match self {
            Target::Raw(off) => Ok(base + off),
            Target::Symbol(sym) => resolve_symbol(sym, base),
            Target::SymbolOffset(sym, off) => {
                let addr = resolve_symbol(sym, base)?;
                Ok(if *off >= 0 {
                    addr + (*off as usize)
                } else {
                    addr - ((-*off) as usize)
                })
            }
            Target::Relative(off) => Ok(if *off >= 0 {
                hook_point + (*off as usize)
            } else {
                hook_point - ((-*off) as usize)
            }),
        }
    }
}

fn resolve_symbol(sym: &str, base: usize) -> Result<usize, PrehookError> {
    use std::ffi::CString;
    let c_sym = CString::new(sym).map_err(|_| PrehookError::SymbolNotFound(sym.to_string()))?;

    let mut addr = unsafe { libc::dlsym(libc::RTLD_DEFAULT, c_sym.as_ptr()) };

    if addr.is_null() {
        addr = unsafe { libc::dlsym(libc::RTLD_NEXT, c_sym.as_ptr()) };
    }

    if addr.is_null() {
        let main_handle = unsafe { libc::dlopen(std::ptr::null(), libc::RTLD_LAZY) };
        if !main_handle.is_null() {
            addr = unsafe { libc::dlsym(main_handle, c_sym.as_ptr()) };
            unsafe { libc::dlclose(main_handle) };
        }
    }

    if addr.is_null() {
        if let Some(offset) = resolve_internal_symbol(sym) {
            return Ok(offset + base);
        }
        return Err(PrehookError::SymbolNotFound(sym.to_string()));
    }

    Ok(addr as usize)
}

/// The variant of the hook being applied.
#[derive(Clone, Debug)]
pub enum HookVariant {
    /// Function hook: interceptions occur at the start of a function,
    /// allowing for original function calls and modified return values.
    Retn,
    /// Inline hook: execution is intercepted, code runs, and jumps back
    /// to the instruction following the hook point.
    JmpBack,
    /// Inline redirection hook: execution is redirected to a dynamic address
    /// returned by the hook.
    JmpToRet,
    /// Static redirection hook: execution is redirected to a fixed
    /// address specified via HookInstance.dest.
    JmpToAddr,
}

/// A thread-safe wrapper for the internal [HookPoint] type from `ilhook`.
///
/// This allows active hooks to be stored in a global static registry.
struct SyncHookPoint(#[allow(dead_code)] HookPoint);

unsafe impl Send for SyncHookPoint {}
unsafe impl Sync for SyncHookPoint {}

/// Represents an individual hook definition.
pub struct HookInstance {
    /// The target specification for the hook point.
    pub target: Target,
    /// The destination for jmp_to_addr hooks (None for other hook types).
    pub dest: Option<Target>,
    /// The kind of hook being applied (Retn, JmpBack, etc.).
    pub hook_variant: HookVariant,
    /// The name of the hook (often the function name it was created from).
    pub name: String,
    /// Callback for Retn hooks.
    pub retn_handler: Option<RetnRoutine>,
    /// Callback for JmpBack hooks.
    pub jmpback_handler: Option<JmpBackRoutine>,
    /// Callback for JmpToRet hooks.
    pub jmptoret_handler: Option<JmpToRetRoutine>,
    /// Callback for JmpToAddr hooks.
    pub jmptoaddr_handler: Option<JmpToAddrRoutine>,
}

// A more aggressive alternative when everything fails
pub fn resolve_internal_symbol(sym_name: &str) -> Option<usize> {
    // 1. Read the memory of the current process directly via /proc/self/exe
    let file = fs::File::open("/proc/self/exe").ok()?;
    let mmap = unsafe { memmap2::Mmap::map(&file).ok()? };
    let elf = Elf::parse(&mmap).ok()?;

    // 2. Iterate over the FULL symbol table (.symtab), not the dynamic one
    for sym in elf.syms.iter() {
        if let Some(name) = elf.strtab.get_at(sym.st_name)
            && name == sym_name
        {
            let sym_value = sym.st_value as usize;
            let load_segment = elf
                .program_headers
                .iter()
                .find(|ph| {
                    ph.p_type == goblin::elf::program_header::PT_LOAD
                        && (sym_value >= ph.p_vaddr as usize
                            && sym_value < (ph.p_vaddr + ph.p_memsz) as usize)
                })
                .expect("Symbol not found in any LOAD segment");

            // st_value is the offset from the base
            return Some(sym_value - load_segment.p_vaddr as usize);
        }
    }
    None
}

impl HookInstance {
    /// Creates a new `Retn` (Function) hook definition.
    pub fn new_retn(target: Target, name: String, handler: RetnRoutine) -> Self {
        Self {
            target,
            dest: None,
            hook_variant: HookVariant::Retn,
            name,
            retn_handler: Some(handler),
            jmpback_handler: None,
            jmptoret_handler: None,
            jmptoaddr_handler: None,
        }
    }

    /// Creates a new `JmpBack` (Inline) hook definition.
    pub fn new_jmpback(target: Target, name: String, handler: JmpBackRoutine) -> Self {
        Self {
            target,
            dest: None,
            hook_variant: HookVariant::JmpBack,
            name,
            retn_handler: None,
            jmpback_handler: Some(handler),
            jmptoret_handler: None,
            jmptoaddr_handler: None,
        }
    }

    /// Creates a new `JmpToRet` (Redirection) hook definition.
    pub fn new_jmptoret(target: Target, name: String, handler: JmpToRetRoutine) -> Self {
        Self {
            target,
            dest: None,
            hook_variant: HookVariant::JmpToRet,
            name,
            retn_handler: None,
            jmpback_handler: None,
            jmptoret_handler: Some(handler),
            jmptoaddr_handler: None,
        }
    }

    /// Creates a new `JmpToAddr` (Static Redirection) hook definition.
    pub fn new_jmptoaddr(
        target: Target,
        dest: Target,
        name: String,
        handler: JmpToAddrRoutine,
    ) -> Self {
        Self {
            target,
            dest: Some(dest),
            hook_variant: HookVariant::JmpToAddr,
            name,
            retn_handler: None,
            jmpback_handler: None,
            jmptoret_handler: None,
            jmptoaddr_handler: Some(handler),
        }
    }

    /// Resolves the target address and applies the hook to the process.
    pub fn apply(&mut self, base: usize) -> Result<HookPoint, PrehookError> {
        let target_addr = self.target.resolve(base, 0)?;

        debug!(
            target: "prehook",
            "Applying hook '{}' at address: 0x{:x} base: 0x{:x}",
            self.name, target_addr, base
        );

        let hook_type = match &self.hook_variant {
            HookVariant::Retn => {
                let h = self.retn_handler.expect("Retn handler missing");
                IlHookType::Retn(h)
            }
            HookVariant::JmpBack => {
                let h = self.jmpback_handler.expect("JmpBack handler missing");
                IlHookType::JmpBack(h)
            }
            HookVariant::JmpToRet => {
                let h = self.jmptoret_handler.expect("JmpToRet handler missing");
                IlHookType::JmpToRet(h)
            }
            HookVariant::JmpToAddr => {
                let h = self.jmptoaddr_handler.expect("JmpToAddr handler missing");
                let dest_addr = self
                    .dest
                    .as_ref()
                    .expect("JmpToAddr requires dest")
                    .resolve(base, target_addr)?;
                IlHookType::JmpToAddr(dest_addr, h)
            }
        };

        let hooker = Hooker::new(
            target_addr,
            hook_type,
            CallbackOption::None,
            0,
            HookFlags::empty(),
        );

        unsafe {
            match hooker.hook() {
                Ok(hp) => {
                    debug!(
                        target: "prehook",
                        "Hook '{}' applied successfully at 0x{:x}",
                        self.name, target_addr
                    );
                    Ok(hp)
                }
                Err(e) => {
                    error!(target: "prehook", "Failed to apply hook '{}': {}", self.name, e);
                    Err(PrehookError::Hook(e.to_string()))
                }
            }
        }
    }
}

/// Global registry for managing hooks.
///
/// Allows hooks to be registered once (typically in a ctor constructor)
/// and applied later (typically in `__libc_start_main`).
#[derive(Default)]
pub struct HookRegistry {
    hooks: OnceLock<Mutex<Vec<HookInstance>>>,
    handles: OnceLock<Mutex<HashMap<String, SyncHookPoint>>>,
}

impl HookRegistry {
    /// Creates a new, empty registry.
    pub const fn new() -> Self {
        Self {
            hooks: OnceLock::new(),
            handles: OnceLock::new(),
        }
    }

    fn get_hooks(&self) -> &Mutex<Vec<HookInstance>> {
        self.hooks.get_or_init(|| Mutex::new(Vec::new()))
    }

    fn get_handles(&self) -> &Mutex<HashMap<String, SyncHookPoint>> {
        self.handles.get_or_init(|| Mutex::new(HashMap::new()))
    }

    /// Adds a hook definition to the registry.
    pub fn register(&self, hook: HookInstance) {
        let mut hooks = self.get_hooks().lock().unwrap();
        hooks.push(hook);
    }

    /// Resolves and applies all registered hooks to the process.
    pub fn apply_all(&self, base: usize) {
        let mut hooks = self.get_hooks().lock().unwrap();
        let mut handles = self.get_handles().lock().unwrap();
        for hook in hooks.iter_mut() {
            match hook.apply(base) {
                Ok(hp) => {
                    handles.insert(hook.name.clone(), SyncHookPoint(hp));
                }
                Err(e) => error!("Failed to apply hook '{}': {}", hook.name, e),
            }
        }
    }

    /// Uninstalls a specific hook by its name.
    pub fn remove_hook(&self, name: &str) {
        let mut handles = self.get_handles().lock().unwrap();
        if handles.remove(name).is_some() {
            debug!(target: "prehook", "Hook '{}' removed and uninstalled", name);
        }
    }

    /// Uninstalls all active hooks in the registry.
    pub fn cleanup_all(&self) {
        let mut handles = self.get_handles().lock().unwrap();
        debug!(target: "prehook", "Cleaning up {} hooks", handles.len());
        // Dropping the handles unhooks them
        handles.clear();
    }
}

/// Global singleton for hook management.
pub static HOOK_REGISTRY: HookRegistry = HookRegistry::new();

#[macro_export]
macro_rules! hook_retn_inner {
    ($offset:expr, $symbol:expr) => {
        #[ctor::ctor]
        fn __prehook_register() {
            let hook = $crate::hook::HookInstance::new_retn(
                $offset,
                $symbol,
                "hook_".to_string(),
                __prehook_handler,
            );
            $crate::hook::HOOK_REGISTRY.register(hook);
        }
    };
}

#[macro_export]
macro_rules! original {
    ($orig_func:expr, $reg:expr) => {{
        let orig: extern "win64" fn(&mut $crate::hook::Registers) -> usize =
            unsafe { std::mem::transmute($orig_func) };
        orig($reg)
    }};
}

extern "win64" fn __prehook_handler(
    _reg: *mut Registers,
    _orig_func_ptr: usize,
    _stack_arg: usize,
) -> usize {
    warn!(target: "prehook", "Hook called but handler not implemented!");
    0
}

#[cfg(test)]
mod tests {
    use super::*;

    extern "win64" fn dummy_retn_handler(
        _reg: *mut Registers,
        _orig: usize,
        _stack: usize,
    ) -> usize {
        0
    }

    unsafe extern "win64" fn dummy_jmpback_handler(_reg: *mut Registers, _user_data: usize) {}

    #[test]
    fn test_hook_variant_enum() {
        let variant = HookVariant::Retn;
        assert!(matches!(variant, HookVariant::Retn));

        let variant = HookVariant::JmpBack;
        assert!(matches!(variant, HookVariant::JmpBack));

        let variant = HookVariant::JmpToRet;
        assert!(matches!(variant, HookVariant::JmpToRet));
    }

    #[test]
    fn test_hook_instance_new_retn() {
        let instance = HookInstance::new_retn(
            Target::Symbol("test_func".to_string()),
            "test_hook".to_string(),
            dummy_retn_handler,
        );

        assert_eq!(instance.target, Target::Symbol("test_func".to_string()));
        assert_eq!(instance.name, "test_hook");
        assert!(matches!(instance.hook_variant, HookVariant::Retn));
        assert!(instance.retn_handler.is_some());
        assert!(instance.jmpback_handler.is_none());
    }

    #[test]
    fn test_hook_instance_new_jmpback() {
        let instance = HookInstance::new_jmpback(
            Target::SymbolOffset("log_hook".to_string(), 0x2000),
            "log_hook".to_string(),
            dummy_jmpback_handler,
        );

        assert_eq!(
            instance.target,
            Target::SymbolOffset("log_hook".to_string(), 0x2000)
        );
        assert!(instance.dest.is_none());
        assert_eq!(instance.name, "log_hook");
        assert!(matches!(instance.hook_variant, HookVariant::JmpBack));
        assert!(instance.jmpback_handler.is_some());
    }

    #[test]
    fn test_hook_registry_new() {
        let registry = HookRegistry::new();
        let hooks = registry.get_hooks().lock().unwrap();
        assert!(hooks.is_empty());
    }

    #[test]
    fn test_hook_registry_register() {
        let registry = HookRegistry::new();
        let instance = HookInstance::new_retn(
            Target::Symbol("test".to_string()),
            "test".to_string(),
            dummy_retn_handler,
        );

        registry.register(instance);

        let hooks = registry.get_hooks().lock().unwrap();
        assert_eq!(hooks.len(), 1);
    }

    #[test]
    fn test_hook_guard_try_enter_first_time() {
        let name = "test_hook";
        assert!(HookGuard::try_enter(name));
        assert!(HookGuard::is_active(name));
        HookGuard::exit(name);
    }

    #[test]
    fn test_hook_guard_prevents_reentry() {
        let name = "recursive_hook";

        let outer = HookGuard::try_enter(name);
        assert!(outer, "First entry should succeed");

        let inner = HookGuard::try_enter(name);
        assert!(!inner, "Second entry should fail (reentry detected)");

        HookGuard::exit(name);

        let after_exit = HookGuard::try_enter(name);
        assert!(after_exit, "Entry should work after exit");
    }

    #[test]
    fn test_hook_guard_exit_removes_name() {
        let name = "exit_test";

        assert!(HookGuard::try_enter(name));
        assert!(HookGuard::is_active(name));

        HookGuard::exit(name);

        assert!(
            !HookGuard::is_active(name),
            "Guard should be inactive after exit"
        );
    }

    #[test]
    fn test_hook_guard_multiple_names() {
        assert!(HookGuard::try_enter("hook_a"));
        assert!(HookGuard::try_enter("hook_b"));
        assert!(HookGuard::is_active("hook_a"));
        assert!(HookGuard::is_active("hook_b"));

        HookGuard::exit("hook_a");
        assert!(!HookGuard::is_active("hook_a"));
        assert!(HookGuard::is_active("hook_b"));

        HookGuard::exit("hook_b");
        assert!(!HookGuard::is_active("hook_b"));
    }
}
