use crate::imagebase::ImageBase;
use crate::structs::*;

pub struct ExportIter<'a> {
    image: &'a ImageBase,
    export_dir: &'a ImageExportDirectory,
    export_index: usize,
}

impl<'a> Iterator for ExportIter<'a> {
    type Item = (usize, &'static [u8]);

    fn next(&mut self) -> Option<Self::Item> {
        match self.export_index < self.export_dir.number_of_name_pointers as usize {
            true => unsafe {
                #[inline(always)]
                unsafe fn u8_nul_terminated(ptr: *const u8) -> &'static [u8] {
                    let mut end = ptr;
                    while *end != 0 { end = end.add(1) }
                    let len = (end as usize) - (ptr as usize);
                    &*core::ptr::slice_from_raw_parts(ptr, len)
                }

                let export_functions = self.image.offset(self.export_dir.export_address_table_rva) as *const u32;
                let export_names = self.image.offset(self.export_dir.name_pointer_rva) as *const u32;
                let export_ordinals = self.image.offset(self.export_dir.ordinal_table_rva) as *const u16;

                let export_name = self.image.offset(*export_names.add(self.export_index));
                let export_ordinal = *export_ordinals.add(self.export_index);
                let export_rva = self.image.offset(*export_functions.add(export_ordinal as usize)) as usize;

                self.export_index += 1;

                Some((export_rva, u8_nul_terminated(export_name)))
            }
            false => None,
        }
    }
}

impl ImageBase {
    pub unsafe fn exports(&self) -> Option<ExportIter> {
        let directory = &self.nt_header().optional_header.data_directory[DIRECTORY_ENTRY_EXPORT];
        if directory.size == 0 || directory.virtual_address == 0 { return None; }
        let export_directory = &*(self.offset(directory.virtual_address) as *const ImageExportDirectory);
        Some(ExportIter {
            image: self,
            export_dir: export_directory,
            export_index: 0,
        })
    }
}