wininskit 0.1.1

use std::ptr;

use windows_sys::Win32::{
    Foundation::{ERROR_FILE_NOT_FOUND, ERROR_SUCCESS},
    System::Registry::{
        HKEY, HKEY_CURRENT_USER, HKEY_LOCAL_MACHINE, KEY_READ, KEY_WOW64_64KEY, KEY_WRITE,
        REG_DWORD, REG_OPTION_NON_VOLATILE, REG_SZ, RegCloseKey, RegCreateKeyExW, RegDeleteTreeW,
        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(())
}

/// 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.
pub fn delete_tree(root: Root, path: &str) -> Result<()> {
    let path = wide(path);
    let status = unsafe { RegDeleteTreeW(root.handle(), path.as_ptr()) };
    if status == ERROR_SUCCESS || status == ERROR_FILE_NOT_FOUND {
        return Ok(());
    }
    Err(Error::code("RegDeleteTreeW", 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) };
    }
}