wininskit 0.1.2

use std::ptr;

use windows_sys::Win32::{
    Foundation::{ERROR_FILE_NOT_FOUND, ERROR_NO_MORE_ITEMS, ERROR_SUCCESS},
    Storage::FileSystem::DELETE,
    System::Registry::{
        HKEY, HKEY_CURRENT_USER, HKEY_LOCAL_MACHINE, KEY_ENUMERATE_SUB_KEYS, KEY_QUERY_VALUE,
        KEY_READ, KEY_SET_VALUE, KEY_WOW64_64KEY, KEY_WRITE, REG_DWORD, REG_OPTION_NON_VOLATILE,
        REG_SZ, RegCloseKey, RegCreateKeyExW, RegDeleteKeyExW, RegDeleteTreeW, RegEnumKeyExW,
        RegOpenKeyExW, RegQueryValueExW, RegSetValueExW,
    },
};

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

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Root {
    LocalMachine,
    CurrentUser,
}

impl Root {
    fn handle(self) -> HKEY {
        match self {
            Root::LocalMachine => HKEY_LOCAL_MACHINE,
            Root::CurrentUser => HKEY_CURRENT_USER,
        }
    }
}

/// Reads a string value, or `None` when either the key or the value is absent.
pub fn read_string(root: Root, path: &str, name: &str) -> Result<Option<String>> {
    let key = match Key::open(root, path, KEY_READ) {
        Ok(key) => key,
        Err(error) if error.is_not_found() => return Ok(None),
        Err(error) => return Err(error),
    };

    let name = wide(name);
    let mut kind = 0u32;
    let mut size = 0u32;

    let status = unsafe {
        RegQueryValueExW(
            key.0,
            name.as_ptr(),
            ptr::null_mut(),
            &mut kind,
            ptr::null_mut(),
            &mut size,
        )
    };
    if status == ERROR_FILE_NOT_FOUND {
        return Ok(None);
    }
    if status != ERROR_SUCCESS {
        return Err(Error::code("RegQueryValueExW", status));
    }

    let mut buffer = vec![0u16; (size as usize).div_ceil(2) + 1];
    let mut size = (buffer.len() * 2) as u32;
    let status = unsafe {
        RegQueryValueExW(
            key.0,
            name.as_ptr(),
            ptr::null_mut(),
            &mut kind,
            buffer.as_mut_ptr().cast(),
            &mut size,
        )
    };
    if status != ERROR_SUCCESS {
        return Err(Error::code("RegQueryValueExW", status));
    }
    Ok(Some(from_wide(&buffer)))
}

pub fn write_string(root: Root, path: &str, name: &str, value: &str) -> Result<()> {
    let key = Key::create(root, path)?;
    let name = wide(name);
    let value = wide(value);
    // The length includes the terminator, which is what lets anything reading
    // it back get a properly terminated string.
    let bytes = (value.len() * 2) as u32;
    let status = unsafe {
        RegSetValueExW(
            key.0,
            name.as_ptr(),
            0,
            REG_SZ,
            value.as_ptr().cast(),
            bytes,
        )
    };
    if status != ERROR_SUCCESS {
        return Err(Error::code("RegSetValueExW", status));
    }
    Ok(())
}

pub fn write_dword(root: Root, path: &str, name: &str, value: u32) -> Result<()> {
    let key = Key::create(root, path)?;
    let name = wide(name);
    let status = unsafe {
        RegSetValueExW(
            key.0,
            name.as_ptr(),
            0,
            REG_DWORD,
            (&raw const value).cast(),
            4,
        )
    };
    if status != ERROR_SUCCESS {
        return Err(Error::code("RegSetValueExW", status));
    }
    Ok(())
}

/// The names of the keys directly under `path`, in the order the registry
/// lists them. A key that is not there has nothing under it.
pub fn subkeys(root: Root, path: &str) -> Result<Vec<String>> {
    let key = match Key::open(root, path, KEY_READ) {
        Ok(key) => key,
        Err(error) if error.is_not_found() => return Ok(Vec::new()),
        Err(error) => return Err(error),
    };

    let mut names = Vec::new();
    for index in 0.. {
        // A key name is at most 255 characters, plus the terminator.
        let mut buffer = [0u16; 256];
        let mut length = buffer.len() as u32;
        let status = unsafe {
            RegEnumKeyExW(
                key.0,
                index,
                buffer.as_mut_ptr(),
                &mut length,
                ptr::null_mut(),
                ptr::null_mut(),
                ptr::null_mut(),
                ptr::null_mut(),
            )
        };
        if status == ERROR_NO_MORE_ITEMS {
            break;
        }
        if status != ERROR_SUCCESS {
            return Err(Error::code("RegEnumKeyExW", status));
        }
        names.push(from_wide(&buffer[..length as usize]));
    }
    Ok(names)
}

/// Removes a key and everything under it. A key that is already gone is not an
/// error, because uninstalling twice should not fail the second time. An
/// empty path names the root itself and is refused.
pub fn delete_tree(root: Root, path: &str) -> Result<()> {
    if path.trim_matches('\\').is_empty() {
        return Err(Error::saying(
            "RegDeleteTreeW",
            "refused to delete everything under a root key",
        ));
    }

    // RegDeleteTreeW has no view flag of its own; it works in whichever view
    // the handle it is given was opened in. So the key is opened in the
    // 64-bit view like every other, emptied through that handle, and then
    // deleted by name in the same view.
    let access = DELETE | KEY_ENUMERATE_SUB_KEYS | KEY_QUERY_VALUE | KEY_SET_VALUE;
    let key = match Key::open(root, path, access) {
        Ok(key) => key,
        Err(error) if error.is_not_found() => return Ok(()),
        Err(error) => return Err(error),
    };
    let status = unsafe { RegDeleteTreeW(key.0, ptr::null()) };
    if status != ERROR_SUCCESS {
        return Err(Error::code("RegDeleteTreeW", status));
    }
    drop(key);

    let path = wide(path);
    let status = unsafe { RegDeleteKeyExW(root.handle(), path.as_ptr(), KEY_WOW64_64KEY, 0) };
    if status == ERROR_SUCCESS || status == ERROR_FILE_NOT_FOUND {
        return Ok(());
    }
    Err(Error::code("RegDeleteKeyExW", status))
}

/// An open key that closes itself.
///
/// Every open asks for the 64-bit view explicitly. The installer is 64-bit so
/// it would get that anyway, but the keys it writes are read by other things
/// and being explicit means one less way for a value to be written somewhere
/// nothing looks.
struct Key(HKEY);

impl Key {
    fn open(root: Root, path: &str, access: u32) -> Result<Self> {
        let path = wide(path);
        let mut key: HKEY = ptr::null_mut();
        let status = unsafe {
            RegOpenKeyExW(
                root.handle(),
                path.as_ptr(),
                0,
                access | KEY_WOW64_64KEY,
                &mut key,
            )
        };
        if status != ERROR_SUCCESS {
            return Err(Error::code("RegOpenKeyExW", status));
        }
        Ok(Self(key))
    }

    fn create(root: Root, path: &str) -> Result<Self> {
        let path = wide(path);
        let mut key: HKEY = ptr::null_mut();
        let status = unsafe {
            RegCreateKeyExW(
                root.handle(),
                path.as_ptr(),
                0,
                ptr::null(),
                REG_OPTION_NON_VOLATILE,
                KEY_WRITE | KEY_WOW64_64KEY,
                ptr::null(),
                &mut key,
                ptr::null_mut(),
            )
        };
        if status != ERROR_SUCCESS {
            return Err(Error::code("RegCreateKeyExW", status));
        }
        Ok(Self(key))
    }
}

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

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

    /// Every Windows has this key, and Uninstall under it.
    #[test]
    fn subkeys_lists_what_is_under_a_key() {
        let under = subkeys(
            Root::LocalMachine,
            r"Software\Microsoft\Windows\CurrentVersion",
        )
        .expect("a key every machine has");
        assert!(under.iter().any(|name| name == "Uninstall"), "{under:?}");
    }

    #[test]
    fn a_key_that_is_not_there_has_nothing_under_it() {
        let under = subkeys(Root::CurrentUser, r"Software\wininskit-never-existed").unwrap();
        assert!(under.is_empty());
    }

    /// An empty path used to reach RegDeleteTreeW with the root handle, which
    /// would have taken everything under HKCU or HKLM.
    #[test]
    fn deleting_a_root_is_refused() {
        for root in ["", r"\"] {
            assert!(delete_tree(Root::CurrentUser, root).is_err(), "{root:?}");
        }
    }

    #[test]
    fn deleting_a_tree_takes_the_key_and_everything_under_it() {
        const KEY: &str = r"Software\wininskit-test-delete-tree";
        write_string(
            Root::CurrentUser,
            &format!(r"{KEY}\Line\Deep"),
            "name",
            "value",
        )
        .unwrap();
        write_string(Root::CurrentUser, KEY, "", "folder").unwrap();
        assert_eq!(
            subkeys(Root::CurrentUser, KEY).unwrap(),
            vec!["Line".to_string()]
        );

        delete_tree(Root::CurrentUser, KEY).expect("delete it");
        assert_eq!(read_string(Root::CurrentUser, KEY, "").unwrap(), None);
        assert!(subkeys(Root::CurrentUser, KEY).unwrap().is_empty());
        // Twice is not a failure.
        delete_tree(Root::CurrentUser, KEY).expect("already gone");
    }
}