//ll_utils.rs
//Low Level cross-platform Memory Utils
use std::arch::asm;
use std::fs::File;
use std::io::{BufRead, BufReader};
use std::ops::Neg;
use thiserror::Error;
use crate::ll_utils::syscalls::{MprotFlags, MPROTECT_EXECUTE, MPROTECT_NONE, MPROTECT_READ, MPROTECT_WRITE};

#[derive(Error, Debug)]
pub enum MemProtectError {
    #[error("Address is either an invalid pointer or not aligned to page boundaries")]
    InvalidPointer,
    #[error("Invalid protection combination")]
    InvalidProtection,
    //ENOMEM https://man7.org/linux/man-pages/man2/mprotect.2.html
    #[error("Protected memory areas are not accessible, likely unmapped memory")]
    InternalError,
}

#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
mod syscalls {
    pub type MprotFlags = u32;

    pub const EINVAL: u64 = 22;
    pub const ENOMEM: u64 = 12;
    pub const SYSCALL_MPROTECT: u32 = 10;
    pub const MPROTECT_READ: MprotFlags = 0x1;
    pub const MPROTECT_WRITE: MprotFlags = 0x2;
    pub const MPROTECT_EXECUTE: MprotFlags = 0x4;
    pub const MPROTECT_NONE: MprotFlags = 0x0;
}

#[cfg(target_os = "linux")]
fn ll_get_memory_protection(addr: usize) -> Option<MprotFlags> {
    let file = File::open("/proc/self/maps").ok()?;
    let reader = BufReader::new(file);

    for line in reader.lines().flatten() {
        // Line format: start-end perms offset dev inode pathname?
        let parts: Vec<&str> = line.split_whitespace().collect();
        if parts.len() < 2 {
            continue;
        }

        let range = parts[0];
        let perms = parts[1];

        let mut range_parts = range.split('-');
        let start = usize::from_str_radix(range_parts.next()?, 16).ok()?;
        let end = usize::from_str_radix(range_parts.next()?, 16).ok()?;

        if addr >= start && addr < end {
            let mut flags = MPROTECT_NONE;

            if perms.contains("x") {
                flags = flags | syscalls::MPROTECT_EXECUTE;
            }

            if perms.contains("r") {
                flags = flags | syscalls::MPROTECT_READ;
            }

            if perms.contains("w") {
                flags = flags | syscalls::MPROTECT_WRITE;
            }
            return Some(flags);
        }
    }

    None
}


#[cfg(all(target_os = "linux", target_arch = "x86_64"))]
fn ll_linux_mprotect(addr: usize, length: usize, prot: syscalls::MprotFlags) -> Result<(), MemProtectError> {

    //addr needs to be page aligned, will need to check how that works easily with no deps

    let syscall_result: isize = 0;

    unsafe {
        //Syscall goooo
        asm!(
        "syscall",
        in("rax") syscalls::SYSCALL_MPROTECT,
        in("rdi") addr,
        in("rsi") length,
        in("rdx") prot,
        lateout("rax") syscall_result,
        clobber: "rcx", "r11", "memory",
        );
    }

    // If there was no error, return
    if syscall_result >= 0 {
        return Ok(());
    }

    //Flip to error code, then convert to u64, this is always safe
    let syscall_result = syscall_result.neg() as u64;

    //We are on the error path!
    let err = match syscall_result {
        syscalls::EINVAL => {
           //This could be either invalid flags or an invalid pointer, let's try to find out
            let possible_flag_combinations_mask = !(syscalls::MPROTECT_READ | syscalls::MPROTECT_WRITE | syscalls::MPROTECT_EXECUTE);
            //If our masked prot flags have any bits set now, we definitely have invalid flags!
            if possible_flag_combinations_mask & prot > 0 {
                return Err(MemProtectError::InvalidProtection)
            }
            //The other possible error is a pointer issue
            Err(MemProtectError::InvalidPointer)
        },
        syscalls::ENOMEM => {
          Err(MemProtectError::InternalError)
        },
        _ => {
            panic!("Unsupported syscall result: {}", syscall_result);
        }
    };
    err

}

#[cfg(target_os = "linux")]
pub fn ll_write_to_mem_internal(addr: usize, data: &[u8]) -> Result<(), MemProtectError> {
    let old_protect = ll_get_memory_protection(addr)
        .unwrap_or(MPROTECT_READ | MPROTECT_EXECUTE);

    //Memory is not writable
    if old_protect & syscalls::MPROTECT_WRITE == 0 {
        ll_linux_mprotect(addr, data.len(), MPROTECT_WRITE)?;
    }

    let dst_ptr = addr as *mut u8;
    let src_ptr = data.as_ptr();

    //I do not think we can make the copy more safe...
    unsafe {
        std::ptr::copy_nonoverlapping(src_ptr, dst_ptr, data.len());
    }


    //If old prot was not already readable, reapply the old options
    if old_protect & syscalls::MPROTECT_WRITE == 0 {
        ll_linux_mprotect(addr, data.len(), old_protect)?;
    }

    Ok(())
}

//vtable.rs
use crate::ll_utils::ll_write_to_mem_internal;
use std::slice;

struct VTableHook<'a, T> {
    original: Option<&'a T>,
    replacement: T,
    base_addr: usize,
    offset: usize,
    installed: bool,
}

impl<T> VTableHook<'_, T> {
    pub fn new(base_addr: usize, replacement: T, offset: usize) -> Self {
        Self {
            original: None,
            replacement,
            offset,
            base_addr,
            installed: false,
        }
    }


    // Is this safeish????
    pub fn install(&mut self) {
        let vtable_addr = self.base_addr + self.offset;
        //This needs to be protected better
        //We read the actual address of the original function from the vtable and safe it
        unsafe {
            let vtable_func_addr = *(vtable_addr as *mut usize);
            self.original = Some(std::mem::transmute::<usize, &T>(vtable_func_addr));
        }

        let bytes: &[u8] = unsafe {
            slice::from_raw_parts(
                &self.replacement as *const T as *const u8,
                size_of::<T>(),
            )
        };

        ll_write_to_mem_internal(vtable_addr, &bytes).expect("Invalid");
    }
}