use sub_pe::imagebase::ImageBase;


#[inline(always)]
pub unsafe fn find_module(hash: u64) -> Option<&'static ImageBase> {
    use sub_core::fnv1::hash_utf16;
    loaded_modules()
        .find(|(_, n)|hash_utf16(*n) == hash)
        .map(|(i,_)|i)
}

#[inline(always)]
pub unsafe fn process_image() -> &'static ImageBase {
    let mut process_exe: *const ImageBase;
    core::arch::asm!(
        "mov {x}, gs:[60h]",    // TEB->PEB
        "mov {x}, [{x} + 10h]", // PEB->ImageBaseAddress
        x = out(reg) process_exe,
    );
    &*process_exe
}

#[inline(always)]
pub unsafe fn loaded_modules() -> ModuleIter {
    let mut module_link: *const LDR_DATA_TABLE_ENTRY;
    core::arch::asm!(
    "mov {x}, gs:[60h]",    // TEB->PEB
    "mov {x}, [{x} + 18h]", // PEB->LDR
    "mov {x}, [{x} + 10h]", // LDR->InLoadOrderModuleList
    x = out(reg) module_link,
    );
    ModuleIter {
        entry: (*module_link).prev,
        head: (*module_link).prev,
    }
}

pub struct ModuleIter {
    entry:  *const LDR_DATA_TABLE_ENTRY,
    head:   *const LDR_DATA_TABLE_ENTRY,
}

impl Iterator for ModuleIter {
    type Item = (&'static ImageBase, &'static [u16]);

    fn next(&mut self) -> Option<Self::Item> {
        unsafe {
            self.entry = (&*self.entry).next;
            match self.entry == self.head {
                true => { None }
                false => {
                    let module = (*self.entry).module;
                    let name = (*self.entry).name.as_slice();
                    Some((&*module,name))
                }
            }
        }
    }
}



/*
    FFI STRUCTS
 */
#[repr(C)]
#[derive(Copy, Clone)]
#[allow(non_camel_case_types)]
struct UNICODE_STRING {
    pub length: u16,
    pub capacity: u16,
    pub buffer: *const u16,
}

#[repr(C)]
#[allow(non_camel_case_types)]
struct LDR_DATA_TABLE_ENTRY {
    /* 0x00 */ pub next:           *const LDR_DATA_TABLE_ENTRY,
    /* 0x08 */ pub prev:           *const LDR_DATA_TABLE_ENTRY,
    /* 0x10 */ pub reserved2:      [usize;4],
    /* 0x30 */ pub module:         *const ImageBase,
    /* 0x38 */ pub entry_point:    *const (),
    /* 0x40 */ pub reserved3:      usize,
    /* 0x48 */ pub path:           UNICODE_STRING,
    /* 0x58 */ pub name:           UNICODE_STRING,
}

impl UNICODE_STRING {
    pub fn as_slice(&self) -> &'static [u16] {
        unsafe { core::slice::from_raw_parts(self.buffer, (self.length / 2) as usize) }
    }
}