windows-eco 0.1.0

use core::fmt;

use bitflags::bitflags;
use windows::Win32::System::Threading::{
    GetCurrentProcess, GetProcessInformation, PROCESS_POWER_THROTTLING_CURRENT_VERSION,
    PROCESS_POWER_THROTTLING_EXECUTION_SPEED, PROCESS_POWER_THROTTLING_IGNORE_TIMER_RESOLUTION,
    SetProcessInformation,
};
use winver::WindowsVersion;

use crate::{PROCESS_POWER_THROTTLING, WIN11_22H2};

bitflags! {
    /// Process Power Throttling Control Mask flags
    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
    pub struct PowerThrottlingControlMask: u32 {
        /// Controls execution speed throttling
        const EXECUTION_SPEED = PROCESS_POWER_THROTTLING_EXECUTION_SPEED;
        /// Controls ignore timer resolution throttling
        const IGNORE_TIMER_RESOLUTION = PROCESS_POWER_THROTTLING_IGNORE_TIMER_RESOLUTION;
    }
}

bitflags! {
    /// Process Power Throttling State Mask flags
    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
    pub struct PowerThrottlingStateMask: u32 {
        /// Execution speed is being throttled
        const EXECUTION_SPEED = 0x1;
        /// Timer resolution requests are being ignored
        const IGNORE_TIMER_RESOLUTION = 0x4;
    }
}

/// Represents the power throttling state of a process.
#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct ProcessPowerThrottlingState {
    /// The version of this structure.
    ///
    /// Must be set to `PROCESS_POWER_THROTTLING_CURRENT_VERSION`
    version: u32,
    /// See [PowerThrottlingControlMask]
    control_mask: u32,
    /// See [PowerThrottlingStateMask]
    state_mask: u32,
}

const _: () = assert!(core::mem::size_of::<ProcessPowerThrottlingState>() == 12);

impl Default for ProcessPowerThrottlingState {
    fn default() -> Self {
        Self {
            version: PROCESS_POWER_THROTTLING_CURRENT_VERSION,
            control_mask: 0,
            state_mask: 0,
        }
    }
}

macro_rules! check_win11_22h2 {
    () => {{
        let ver = WindowsVersion::from_ntdll_dll()?;
        if ver.build < WIN11_22H2 {
            return Err(crate::Error::NotAvailable(ver.build));
        }
    }};
}

/// Retrieves the current power throttling state of the calling process.
pub(crate) fn get_power_throttling_state(
    version: u32,
) -> Result<ProcessPowerThrottlingState, crate::Error> {
    check_win11_22h2!();

    let mut process_power_throttling = ProcessPowerThrottlingState {
        version,
        ..Default::default()
    };
    let process_information_size = core::mem::size_of::<ProcessPowerThrottlingState>() as u32;

    unsafe {
        GetProcessInformation(
            GetCurrentProcess(),
            PROCESS_POWER_THROTTLING,
            &mut process_power_throttling as *mut _ as *mut std::ffi::c_void,
            process_information_size,
        )?;
    }

    Ok(process_power_throttling)
}

/// Sets the current power throttling state of the calling process.
pub(crate) fn set_power_throttling_state(
    version: u32,
    state_mask: u32,
    control_mask: u32,
) -> Result<(), crate::Error> {
    check_win11_22h2!();

    let mut process_power_throttling = ProcessPowerThrottlingState {
        version,
        state_mask,
        control_mask,
        ..Default::default()
    };
    let process_information_size = core::mem::size_of::<ProcessPowerThrottlingState>() as u32;

    unsafe {
        SetProcessInformation(
            GetCurrentProcess(),
            PROCESS_POWER_THROTTLING,
            &mut process_power_throttling as *mut _ as *mut std::ffi::c_void,
            process_information_size,
        )?;
    }

    Ok(())
}

impl ProcessPowerThrottlingState {
    /// Creates a new [ProcessPowerThrottlingState] with specified control and state masks
    pub fn new(
        control_mask: PowerThrottlingControlMask,
        state_mask: PowerThrottlingStateMask,
    ) -> Self {
        Self {
            control_mask: control_mask.bits(),
            state_mask: state_mask.bits(),
            ..Default::default()
        }
    }

    /// Create [ProcessPowerThrottlingState] from current process
    pub fn from_windows() -> Result<Self, crate::Error> {
        get_power_throttling_state(PROCESS_POWER_THROTTLING_CURRENT_VERSION)
    }

    /// Applies this power throttling state to the current process
    pub fn apply(&self) -> Result<(), crate::Error> {
        set_power_throttling_state(self.version, self.state_mask, self.control_mask)
    }

    /// Creates a ProcessPowerThrottlingState that enables execution speed throttling
    pub fn enable_execution_speed_throttling(&mut self) {
        self.control_mask |= PowerThrottlingControlMask::EXECUTION_SPEED.bits();
        self.state_mask |= PowerThrottlingStateMask::EXECUTION_SPEED.bits();
    }

    /// Creates a ProcessPowerThrottlingState that disables execution speed throttling
    pub fn disable_execution_speed_throttling(&mut self) {
        self.control_mask |= PowerThrottlingControlMask::EXECUTION_SPEED.bits();
        self.state_mask |= PowerThrottlingStateMask::empty().bits();
    }

    /// Creates a ProcessPowerThrottlingState that enables timer resolution throttling
    pub fn enable_timer_resolution_throttling(&mut self) {
        self.control_mask |= PowerThrottlingControlMask::IGNORE_TIMER_RESOLUTION.bits();
        self.state_mask |= PowerThrottlingStateMask::IGNORE_TIMER_RESOLUTION.bits();
    }

    /// Creates a ProcessPowerThrottlingState that disables timer resolution throttling
    pub fn disable_timer_resolution_throttling(&mut self) {
        self.control_mask |= PowerThrottlingControlMask::IGNORE_TIMER_RESOLUTION.bits();
        self.state_mask |= PowerThrottlingStateMask::empty().bits();
    }

    /// Creates a ProcessPowerThrottlingState that enables all available throttling
    pub fn enable_all_throttling(&mut self) {
        self.control_mask |= PowerThrottlingControlMask::all().bits();
        self.state_mask |= PowerThrottlingStateMask::all().bits();
    }

    /// Creates a ProcessPowerThrottlingState that disables all throttling
    pub fn disable_all_throttling(&mut self) {
        self.control_mask |= PowerThrottlingControlMask::all().bits();
        self.state_mask |= PowerThrottlingStateMask::empty().bits();
    }

    /// Get control mask as bitflags
    pub fn control_flags(&self) -> PowerThrottlingControlMask {
        PowerThrottlingControlMask::from_bits_truncate(self.control_mask)
    }

    /// Get state mask as bitflags
    pub fn state_flags(&self) -> PowerThrottlingStateMask {
        PowerThrottlingStateMask::from_bits_truncate(self.state_mask)
    }

    /// Check if execution speed throttling is enabled
    pub fn is_execution_speed_throttled(&self) -> bool {
        self.state_flags()
            .contains(PowerThrottlingStateMask::EXECUTION_SPEED)
    }

    /// Check if timer resolution throttling is enabled
    pub fn is_ignore_timer_resolution_throttled(&self) -> bool {
        self.state_flags()
            .contains(PowerThrottlingStateMask::IGNORE_TIMER_RESOLUTION)
    }

    /// Check if execution speed control is enabled
    pub fn is_execution_speed_controlled(&self) -> bool {
        self.control_flags()
            .contains(PowerThrottlingControlMask::EXECUTION_SPEED)
    }

    /// Check if timer resolution control is enabled
    pub fn is_timer_resolution_controlled(&self) -> bool {
        self.control_flags()
            .contains(PowerThrottlingControlMask::IGNORE_TIMER_RESOLUTION)
    }

    /// Retrieve the underlying version of power throttling API
    pub fn version(&self) -> u32 {
        self.version
    }

    /// Sets the underlying version of power throttling API
    pub fn set_version(&mut self, version: u32) {
        self.version = version;
    }
}

impl fmt::Display for ProcessPowerThrottlingState {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        let state_flags = self.state_flags();
        let control_flags = self.control_flags();

        if state_flags.is_empty() {
            write!(f, "No Throttling Active")
        } else {
            let mut descriptions = Vec::new();

            if state_flags.contains(PowerThrottlingStateMask::EXECUTION_SPEED) {
                descriptions.push("Execution Speed Throttled");
            }

            if state_flags.contains(PowerThrottlingStateMask::IGNORE_TIMER_RESOLUTION) {
                descriptions.push("Timer Resolution Throttled");
            }

            write!(f, "{}", descriptions.join(", "))?;

            // Show controlled flags if different from state
            if !control_flags.is_empty()
                && control_flags
                    != PowerThrottlingControlMask::from_bits_truncate(state_flags.bits())
            {
                write!(f, " (Controlled: {})", control_flags)?;
            }

            Ok(())
        }
    }
}

impl fmt::Display for PowerThrottlingControlMask {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if self.is_empty() {
            write!(f, "None")
        } else {
            let mut flags = Vec::new();
            if self.contains(Self::EXECUTION_SPEED) {
                flags.push("ExecutionSpeed");
            }
            if self.contains(Self::IGNORE_TIMER_RESOLUTION) {
                flags.push("TimerResolution");
            }
            write!(f, "{}", flags.join("|"))
        }
    }
}

impl fmt::Display for PowerThrottlingStateMask {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if self.is_empty() {
            write!(f, "None")
        } else {
            let mut flags = Vec::new();
            if self.contains(Self::EXECUTION_SPEED) {
                flags.push("ExecutionSpeed");
            }
            if self.contains(Self::IGNORE_TIMER_RESOLUTION) {
                flags.push("TimerResolution");
            }
            write!(f, "{}", flags.join("|"))
        }
    }
}