winkit 0.1.0

use std::{mem, path::Path};

use std::time::Duration;

use windows_sys::Win32::{
    Foundation::{CloseHandle, INVALID_HANDLE_VALUE, MAX_PATH, WAIT_OBJECT_0},
    System::{
        Diagnostics::ToolHelp::{
            CreateToolhelp32Snapshot, PROCESSENTRY32W, Process32FirstW, Process32NextW,
            TH32CS_SNAPPROCESS,
        },
        Threading::{
            OpenProcess, PROCESS_QUERY_LIMITED_INFORMATION, PROCESS_TERMINATE,
            QueryFullProcessImageNameW, TerminateProcess, WaitForSingleObject,
        },
    },
};

use crate::{
    error::{Error, Result},
    wide::from_wide,
};

#[derive(Debug, Clone)]
pub struct Process {
    pub id: u32,
    pub name: String,
    /// The executable on disk, when it could be read. A process owned by
    /// another user can refuse this even to an administrator.
    pub image: Option<String>,
}

/// Ends a process and waits for it to actually be gone.
///
/// This is a kill, not a request: it does not ask the program to save anything
/// or let it refuse, so nothing should call it without the person having said
/// so first. A process that has already exited counts as closed, since that is
/// the state the caller wanted.
pub fn close(id: u32, wait: Duration) -> Result<()> {
    // The standard right to wait on a handle, without which the wait below
    // fails rather than waits. windows-sys files it under storage access
    // rights, which is not where a process handle would think to look.
    const SYNCHRONIZE: u32 = 0x0010_0000;

    unsafe {
        let handle = OpenProcess(
            PROCESS_TERMINATE | PROCESS_QUERY_LIMITED_INFORMATION | SYNCHRONIZE,
            0,
            id,
        );
        if handle.is_null() {
            // Gone between being listed and being closed, which is a fine way
            // for it to end.
            return if running(id) {
                Err(Error::last("OpenProcess"))
            } else {
                Ok(())
            };
        }
        let handle = Handle(handle);

        if TerminateProcess(handle.0, 1) == 0 && running(id) {
            return Err(Error::last("TerminateProcess"));
        }
        // Terminating asks the kernel to end it, and the handles it held are
        // only released once it has, which is what the caller is waiting for.
        let waited = WaitForSingleObject(handle.0, wait.as_millis() as u32);
        if waited != WAIT_OBJECT_0 {
            return Err(Error::saying(
                "WaitForSingleObject",
                format!("{id} was ended but has not gone yet"),
            ));
        }
    }
    Ok(())
}

/// Whether a process with this id is still listed.
fn running(id: u32) -> bool {
    snapshot().is_ok_and(|processes| processes.iter().any(|process| process.id == id))
}

/// Closes a handle when it goes out of scope, including on the error paths.
struct Handle(windows_sys::Win32::Foundation::HANDLE);

impl Drop for Handle {
    fn drop(&mut self) {
        unsafe { CloseHandle(self.0) };
    }
}

/// Every running process whose executable lives under `directory`.
///
/// This is the question worth asking before replacing an install: not "is
/// anything called UnrealEditor.exe running" but "is anything running out of
/// the folder I am about to overwrite".
pub fn running_from(directory: &Path) -> Result<Vec<Process>> {
    let prefix = normalize(&directory.to_string_lossy());
    let mut found = Vec::new();
    for process in snapshot()? {
        let Some(image) = process.image.as_deref() else {
            continue;
        };
        if normalize(image).starts_with(&prefix) {
            found.push(process);
        }
    }
    Ok(found)
}

/// Lower-cased, backslash-separated and with the verbatim prefix removed, so
/// two spellings of the same directory compare equal.
fn normalize(path: &str) -> String {
    let path = path.strip_prefix(r"\\?\").unwrap_or(path);
    let mut text = path.replace('/', "\\").to_lowercase();
    if !text.ends_with('\\') {
        text.push('\\');
    }
    text
}

fn snapshot() -> Result<Vec<Process>> {
    let mut processes = Vec::new();
    unsafe {
        let snapshot = CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0);
        if snapshot == INVALID_HANDLE_VALUE {
            return Err(Error::last("CreateToolhelp32Snapshot"));
        }
        let mut entry: PROCESSENTRY32W = mem::zeroed();
        entry.dwSize = mem::size_of::<PROCESSENTRY32W>() as u32;

        let mut ok = Process32FirstW(snapshot, &mut entry);
        while ok != 0 {
            let name = from_wide(&entry.szExeFile);
            if !name.is_empty() {
                processes.push(Process {
                    id: entry.th32ProcessID,
                    image: image_of(entry.th32ProcessID),
                    name,
                });
            }
            ok = Process32NextW(snapshot, &mut entry);
        }
        CloseHandle(snapshot);
    }
    Ok(processes)
}

/// The full path of a running process, or `None` when it cannot be read. A
/// refusal here is normal - system processes do not hand theirs over - so it is
/// not treated as a failure of the whole enumeration.
fn image_of(id: u32) -> Option<String> {
    unsafe {
        let handle = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, 0, id);
        if handle.is_null() {
            return None;
        }
        let mut buffer = [0u16; MAX_PATH as usize];
        let mut length = buffer.len() as u32;
        let ok = QueryFullProcessImageNameW(handle, 0, buffer.as_mut_ptr(), &mut length);
        CloseHandle(handle);
        if ok == 0 {
            return None;
        }
        Some(from_wide(&buffer[..length as usize]))
    }
}

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

    #[test]
    fn closing_ends_the_process_and_waits_for_it() {
        let mut child = std::process::Command::new("cmd")
            .args(["/c", "pause"])
            .stdin(std::process::Stdio::piped())
            .stdout(std::process::Stdio::piped())
            .spawn()
            .expect("start something to close");
        let id = child.id();
        assert!(running(id), "it should be listed before being closed");

        close(id, Duration::from_secs(10)).expect("close it");
        assert!(!running(id), "it should be gone once close returns");
        let _ = child.wait();
    }

    #[test]
    fn closing_something_already_gone_is_not_a_failure() {
        let mut child = std::process::Command::new("cmd")
            .args(["/c", "exit"])
            .spawn()
            .expect("start something short");
        let id = child.id();
        let _ = child.wait();
        close(id, Duration::from_secs(5)).expect("already gone counts as closed");
    }
}