use crate::process::ModuleEntry;
use crate::winapi::um::tlhelp32::{
    CreateToolhelp32Snapshot, Module32First, Module32Next, MODULEENTRY32, TH32CS_SNAPMODULE,
    TH32CS_SNAPMODULE32,
};
use std::mem::zeroed;
use crate::winapi::um::winnt::HANDLE;
use crate::winapi::um::handleapi::CloseHandle;

#[allow(missing_debug_implementations)]
pub struct ModuleIterator {
    h_snap: HANDLE,
    entry: MODULEENTRY32,
    end: i32,
}

impl Iterator for ModuleIterator {
    type Item = ModuleEntry;

    fn next(&mut self) -> Option<Self::Item> {
        unsafe {
            if self.end == -1 {
                self.end = Module32First(self.h_snap, &mut self.entry);
                Some(ModuleEntry {
                    module_base: 0,
                    exe_path: crate::to_c_str!(self.entry.szExePath).to_string(),
                    module_size: 0,
                })
            } else if self.end != 0 {
                let entry = Some(ModuleEntry {
                    module_base: self.entry.modBaseAddr as u64,
                    exe_path: crate::to_c_str!(self.entry.szExePath).to_string(),
                    module_size: self.entry.modBaseSize,
                });

                self.end = Module32Next(self.h_snap, &mut self.entry);
                entry
            } else {
                CloseHandle(self.h_snap);
                None
            }
        }
    }
}

impl ModuleIterator {
    pub fn of(process_id: u32) -> Self {
        unsafe {
            let mut entry: MODULEENTRY32 = zeroed();
            entry.dwSize = crate::size_of!(@ entry) as u32;

            Self {
                h_snap: CreateToolhelp32Snapshot(
                    TH32CS_SNAPMODULE32 | TH32CS_SNAPMODULE,
                    process_id,
                ),
                end: 0,
                entry,
            }
        }
    }
}
