wininskit 0.1.2

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

use std::time::Duration;

use windows_sys::Win32::{
    Foundation::{
        CloseHandle, ERROR_INSUFFICIENT_BUFFER, ERROR_INVALID_PARAMETER, GetLastError,
        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() {
            // Read before anything else runs: listing the processes below
            // ends with an error code of its own and would replace this one.
            let failed = Error::last("OpenProcess");
            // No process by that id, or one gone between being listed and
            // being closed, is what the caller wanted.
            return if failed.code == ERROR_INVALID_PARAMETER || !running(id) {
                Ok(())
            } else {
                Err(failed)
            };
        }
        let handle = Handle(handle);

        if TerminateProcess(handle.0, 1) == 0 {
            let failed = Error::last("TerminateProcess");
            // Ending a process that has already ended is refused. The handle
            // says whether it has, where the process list would still show
            // one whose handles are open.
            if WaitForSingleObject(handle.0, 0) != WAIT_OBJECT_0 {
                return Err(failed);
            }
        }
        // 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.
        // A wait too long for the call is cut to the longest it takes: one
        // less than u32::MAX, which would mean forever.
        let millis = u32::try_from(wait.as_millis()).unwrap_or(u32::MAX - 1);
        let waited = WaitForSingleObject(handle.0, millis);
        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> {
    /// The most characters a path can have once long paths are allowed.
    const LONG_PATH: usize = 32_768;

    unsafe {
        let handle = OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, 0, id);
        if handle.is_null() {
            return None;
        }
        let handle = Handle(handle);
        // Nearly every path fits the short buffer. One that does not is
        // reported as too small rather than cut short, and is asked for again
        // with room for the longest path there is.
        for capacity in [MAX_PATH as usize, LONG_PATH] {
            let mut buffer = vec![0u16; capacity];
            let mut length = buffer.len() as u32;
            if QueryFullProcessImageNameW(handle.0, 0, buffer.as_mut_ptr(), &mut length) != 0 {
                return Some(from_wide(&buffer[..length as usize]));
            }
            if GetLastError() != ERROR_INSUFFICIENT_BUFFER {
                return None;
            }
        }
        None
    }
}

#[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");
        // The process list can still show it while this handle is open, so
        // the handle is what says whether it has ended.
        let ended = child.try_wait().expect("ask after it");
        assert!(ended.is_some(), "it should have ended once close returns");
    }

    /// The handle held here keeps the ended process in the process list, which
    /// is exactly the state a process that was listed and then closed by
    /// someone else is in.
    #[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");
        drop(child);
        close(id, Duration::from_secs(5)).expect("no longer listed counts as closed too");
    }
}