joularcore 0.1.0

/*
 * Copyright (c) 2025-2026, Adel Noureddine.
 * All rights reserved. This program and the accompanying materials
 * are made available under the terms of the
 * GNU Lesser General Public License v3.0 only (LGPL-3.0-only)
 * which accompanies this distribution, and is available at
 * https://www.gnu.org/licenses/lgpl-3.0.en.html
 *
 * Author : Adel Noureddine
 */
use crate::cpu::{AppCPUUtilization, CPUUtilization};
use crate::energy::{CPUEnergy, GPUEnergy, PlatformEnergy};
use crate::platform::amdgpu;
use crate::platform::nvidia;
use std::collections::{HashMap, HashSet};
use std::io;
use std::mem;
use std::sync::Mutex;
use std::time::Instant;

use windows::Win32::Foundation::{CloseHandle, FILETIME, HANDLE, INVALID_HANDLE_VALUE};
use windows::Win32::Storage::FileSystem::*;
use windows::Win32::System::Diagnostics::ToolHelp::{
    CreateToolhelp32Snapshot, PROCESSENTRY32W, Process32FirstW, Process32NextW, TH32CS_SNAPPROCESS,
};
use windows::Win32::System::IO::DeviceIoControl;
use windows::Win32::System::Threading::{
    GetProcessTimes, GetSystemTimes, OpenProcess, PROCESS_QUERY_INFORMATION, PROCESS_VM_READ,
};
use windows::core::*;

pub struct WindowsPlatform;

impl WindowsPlatform {
    pub fn new() -> Self {
        WindowsPlatform
    }
}

impl Default for WindowsPlatform {
    fn default() -> Self {
        Self::new()
    }
}

impl PlatformEnergy for WindowsPlatform {
    fn cpu(&self) -> Box<dyn CPUEnergy> {
        Box::new(WindowsCpu::new())
    }

    fn gpu(&self) -> Box<dyn GPUEnergy> {
        Box::new(WindowsGpu)
    }

    fn cpu_usage(&self) -> Box<dyn CPUUtilization> {
        Box::new(WindowsCpuUsage::new())
    }

    fn process_cpu_usage(&self) -> Option<Box<dyn crate::cpu::ProcessCPUUtilization>> {
        Some(Box::new(WindowsProcessUtil::new()))
    }

    fn app_cpu_usage(
        &self,
        refresh_interval: std::time::Duration,
    ) -> Option<Box<dyn AppCPUUtilization>> {
        Some(Box::new(WindowsAppMonitor::new(refresh_interval)))
    }
}

pub struct CpuState {
    prev_idle: u64,
    prev_kernel: u64,
    prev_user: u64,
}

impl CpuState {
    fn new() -> Self {
        Self {
            prev_idle: 0,
            prev_kernel: 0,
            prev_user: 0,
        }
    }
}

pub struct WindowsCpuUsage {
    state: Mutex<CpuState>,
}

impl WindowsCpuUsage {
    pub fn new() -> Self {
        Self {
            state: Mutex::new(CpuState::new()),
        }
    }

    fn filetime_to_u64(ft: &FILETIME) -> u64 {
        ((ft.dwHighDateTime as u64) << 32) | (ft.dwLowDateTime as u64)
    }
}

impl Default for WindowsCpuUsage {
    fn default() -> Self {
        Self::new()
    }
}

impl CPUUtilization for WindowsCpuUsage {
    fn get_cpu_utilization(&self) -> f64 {
        unsafe {
            let mut idle = mem::zeroed();
            let mut kernel = mem::zeroed();
            let mut user = mem::zeroed();

            if GetSystemTimes(Some(&mut idle), Some(&mut kernel), Some(&mut user)).is_err() {
                return 0.0;
            }

            let idle_time = Self::filetime_to_u64(&idle);
            let kernel_time = Self::filetime_to_u64(&kernel);
            let user_time = Self::filetime_to_u64(&user);

            let mut state = self.state.lock().unwrap_or_else(|e| e.into_inner());

            if state.prev_idle == 0 {
                // First call, just store values
                state.prev_idle = idle_time;
                state.prev_kernel = kernel_time;
                state.prev_user = user_time;
                return 0.0;
            }

            let idle_delta = idle_time.saturating_sub(state.prev_idle);
            let kernel_delta = kernel_time.saturating_sub(state.prev_kernel);
            let user_delta = user_time.saturating_sub(state.prev_user);

            let total_delta = kernel_delta + user_delta;
            let usage = if total_delta > 0 {
                (total_delta - idle_delta) as f64 / total_delta as f64
            } else {
                return 0.0;
            };

            state.prev_idle = idle_time;
            state.prev_kernel = kernel_time;
            state.prev_user = user_time;

            usage
        }
    }
}

pub struct WindowsGpu;
impl GPUEnergy for WindowsGpu {
    fn get_power(&self) -> f64 {
        let mut gpu_energy = 0.0;
        if nvidia::is_nvidia_supported() {
            gpu_energy += nvidia::get_nvidia_power();
        }
        if amdgpu::is_amdgpu_supported() {
            gpu_energy += amdgpu::get_amdgpu_power();
        }
        gpu_energy
    }
}

pub struct WindowsCpu {
    rapl: Mutex<Option<RaplDriver>>,
}

impl WindowsCpu {
    pub fn new() -> Self {
        match RaplDriver::new() {
            Ok(rapl) => Self {
                rapl: Mutex::new(Some(rapl)),
            },
            Err(e) => {
                crate::logging::print_warning(&format!(
                    "Scaphandre RAPL driver unavailable: {}. Continuing without CPU power data",
                    e
                ));
                Self {
                    rapl: Mutex::new(None),
                }
            }
        }
    }
}

impl Default for WindowsCpu {
    fn default() -> Self {
        Self::new()
    }
}

impl CPUEnergy for WindowsCpu {
    fn get_power(&self) -> f64 {
        let mut rapl = self.rapl.lock().unwrap_or_else(|e| e.into_inner());
        match rapl.as_mut() {
            Some(driver) => driver.calculate_power().unwrap_or(0.0),
            None => 0.0,
        }
    }
}

// ============================================================================
// Windows Process CPU Utilization
// ============================================================================

use crate::cpu::ProcessCPUUtilization;

/// Windows-based process CPU utilization tracker
pub struct WindowsProcessUtil {
    before_system_time: u64,
    before_process_time: u64,
}

impl WindowsProcessUtil {
    pub fn new() -> Self {
        Self {
            before_system_time: 0,
            before_process_time: 0,
        }
    }

    /// Convert FILETIME to u64 (100-nanosecond intervals)
    fn filetime_to_u64(ft: &FILETIME) -> u64 {
        ((ft.dwHighDateTime as u64) << 32) | (ft.dwLowDateTime as u64)
    }

    /// Get total system CPU time (kernel + user)
    fn get_system_time() -> Option<u64> {
        unsafe {
            let mut idle = mem::zeroed();
            let mut kernel = mem::zeroed();
            let mut user = mem::zeroed();

            if GetSystemTimes(Some(&mut idle), Some(&mut kernel), Some(&mut user)).is_err() {
                return None;
            }

            Some(Self::filetime_to_u64(&kernel) + Self::filetime_to_u64(&user))
        }
    }

    /// Get CPU time used by a process (kernel + user)
    fn get_process_time(pid: u32) -> Option<u64> {
        unsafe {
            let handle = match OpenProcess(PROCESS_QUERY_INFORMATION | PROCESS_VM_READ, false, pid)
            {
                Ok(h) => h,
                Err(_) => return None,
            };

            let mut creation = mem::zeroed();
            let mut exit = mem::zeroed();
            let mut kernel = mem::zeroed();
            let mut user = mem::zeroed();

            let success = GetProcessTimes(handle, &mut creation, &mut exit, &mut kernel, &mut user);
            let _ = CloseHandle(handle);

            if success.is_err() {
                return None;
            }

            Some(Self::filetime_to_u64(&kernel) + Self::filetime_to_u64(&user))
        }
    }
}

impl Default for WindowsProcessUtil {
    fn default() -> Self {
        Self::new()
    }
}

impl ProcessCPUUtilization for WindowsProcessUtil {
    fn get_process_cpu_utilization(&mut self, pid: u32) -> f64 {
        let system_time = match Self::get_system_time() {
            Some(t) => t,
            None => return 0.0,
        };

        let process_time = match Self::get_process_time(pid) {
            Some(t) => t,
            None => return 0.0,
        };

        if self.before_system_time == 0 {
            self.before_system_time = system_time;
            self.before_process_time = process_time;
            return 0.0;
        }

        let system_delta = system_time.saturating_sub(self.before_system_time);
        let process_delta = process_time.saturating_sub(self.before_process_time);

        self.before_system_time = system_time;
        self.before_process_time = process_time;

        if system_delta == 0 {
            0.0
        } else {
            process_delta as f64 / system_delta as f64
        }
    }
}

// ============================================================================
// Windows Application Monitoring
// ============================================================================

/// Windows-specific application monitor using CreateToolhelp32Snapshot
pub struct WindowsAppMonitor {
    process_trackers: HashMap<u32, WindowsProcessUtil>,
    last_pid_sweep: Option<std::time::Instant>,
    cached_pids: Vec<u32>,
    sweep_interval: std::time::Duration,
}

impl WindowsAppMonitor {
    pub fn new(sweep_interval: std::time::Duration) -> Self {
        Self {
            process_trackers: HashMap::new(),
            last_pid_sweep: None,
            cached_pids: Vec::new(),
            sweep_interval,
        }
    }

    /// Get all PIDs for processes matching the given name on Windows
    fn get_pids_by_name(&mut self, app_name: &str) -> Vec<u32> {
        if self.sweep_interval.as_secs() > 0
            && let Some(last) = self.last_pid_sweep
            && last.elapsed() < self.sweep_interval
        {
            return self.cached_pids.clone();
        }

        unsafe {
            let snapshot = match CreateToolhelp32Snapshot(TH32CS_SNAPPROCESS, 0) {
                Ok(h) => h,
                Err(_) => return Vec::new(),
            };

            let mut pids = Vec::new();
            let mut entry: PROCESSENTRY32W = mem::zeroed();
            entry.dwSize = mem::size_of::<PROCESSENTRY32W>() as u32;

            if Process32FirstW(snapshot, &mut entry).is_ok() {
                loop {
                    // Convert process name from wide char to string
                    let process_name = String::from_utf16_lossy(
                        &entry
                            .szExeFile
                            .iter()
                            .take_while(|&&c| c != 0)
                            .copied()
                            .collect::<Vec<u16>>(),
                    );

                    // Match process name (case-insensitive, with or without .exe)
                    let process_name_lower = process_name.to_lowercase();
                    let app_name_lower = app_name.to_lowercase();

                    if process_name_lower == app_name_lower
                        || process_name_lower == format!("{}.exe", app_name_lower)
                        || process_name_lower.trim_end_matches(".exe") == app_name_lower
                    {
                        pids.push(entry.th32ProcessID);
                    }

                    if Process32NextW(snapshot, &mut entry).is_err() {
                        break;
                    }
                }
            }

            let _ = CloseHandle(snapshot);

            if self.sweep_interval.as_secs() > 0 {
                self.cached_pids = pids.clone();
                self.last_pid_sweep = Some(std::time::Instant::now());
            }

            pids
        }
    }

    /// Update process trackers based on current PIDs
    fn update_trackers(&mut self, current_pids: &[u32]) {
        let pid_set: HashSet<u32> = current_pids.iter().copied().collect();
        // Remove trackers for PIDs that no longer exist
        self.process_trackers.retain(|pid, _| pid_set.contains(pid));

        // Add trackers for new PIDs
        for &pid in current_pids {
            self.process_trackers.entry(pid).or_default();
        }
    }
}

impl AppCPUUtilization for WindowsAppMonitor {
    fn set_refresh_interval(&mut self, interval: std::time::Duration) {
        self.sweep_interval = interval;
    }

    fn get_app_cpu_utilization(&mut self, app_name: &str) -> f64 {
        let current_pids = self.get_pids_by_name(app_name);

        if current_pids.is_empty() {
            self.process_trackers.clear();
            return 0.0;
        }

        self.update_trackers(&current_pids);

        let mut total_utilization = 0.0;
        for &pid in &current_pids {
            if let Some(tracker) = self.process_trackers.get_mut(&pid) {
                total_utilization += tracker.get_process_cpu_utilization(pid);
            }
        }

        total_utilization
    }

    fn get_app_pids(&mut self, app_name: &str) -> Vec<u32> {
        self.get_pids_by_name(app_name)
    }
}

// ============================================================================
// RAPL Driver (existing code)
// ============================================================================

// RAPL MSR addresses for Intel
const MSR_INTEL_RAPL_POWER_UNIT: u64 = 0x606;
const MSR_INTEL_PKG_ENERGY_STATUS: u64 = 0x611;

// RAPL MSR addresses for AMD
const MSR_AMD_RAPL_POWER_UNIT: u64 = 0xc0010299;
const MSR_AMD_PKG_ENERGY_STATUS: u64 = 0xc001029b;

// CTL_CODE macro implementation
const fn ctl_code(device_type: u32, function: u32, method: u32, access: u32) -> u32 {
    (device_type << 16) | (access << 14) | (function << 2) | method
}

const FILE_DEVICE_UNKNOWN: u32 = 0x00000022;
const METHOD_BUFFERED: u32 = 0;
const FILE_READ_DATA: u32 = 0x0001;
const FILE_WRITE_DATA: u32 = 0x0002;

#[derive(Debug, Clone, Copy)]
enum CpuVendor {
    Intel,
    Amd,
}

pub struct RaplDriver {
    handle: HANDLE,
    power_unit: f64,
    energy_unit: f64,
    time_unit: f64,
    vendor: CpuVendor,
    pkg: bool,
    last_energy: f64,
    last_read_at: Option<Instant>,
}

impl RaplDriver {
    pub fn new() -> Result<Self> {
        let driver_path = w!("\\\\.\\ScaphandreDriver");

        let handle = unsafe {
            CreateFileW(
                driver_path,
                FILE_GENERIC_READ.0 | FILE_GENERIC_WRITE.0,
                FILE_SHARE_READ | FILE_SHARE_WRITE,
                None,
                OPEN_EXISTING,
                FILE_FLAG_OVERLAPPED,
                None,
            )?
        };

        if handle == INVALID_HANDLE_VALUE {
            return Err(Error::from_thread());
        }

        let mut driver = RaplDriver {
            handle,
            power_unit: 0.0,
            energy_unit: 0.0,
            time_unit: 0.0,
            vendor: CpuVendor::Intel, // Default to Intel, will be detected
            pkg: false,
            last_energy: 0.0,
            last_read_at: None,
        };

        driver.detect_cpu_vendor()?;
        driver.get_energy_units()?;
        driver.check_supported_platform()?;

        Ok(driver)
    }

    fn get_rapl_ctl_code(&self, msr: u64) -> u32 {
        ctl_code(
            FILE_DEVICE_UNKNOWN,
            (msr & 0xFFF) as u32,
            METHOD_BUFFERED,
            FILE_READ_DATA | FILE_WRITE_DATA,
        )
    }

    fn get_data_from_driver(&self, msr: u64) -> Result<u64> {
        let mut reply_data: u64 = 0;
        let mut bytes_returned: u32 = 0;

        let ctl_code = self.get_rapl_ctl_code(msr);

        unsafe {
            DeviceIoControl(
                self.handle,
                ctl_code,
                Some(&msr as *const u64 as *const _),
                std::mem::size_of::<u64>() as u32,
                Some(&mut reply_data as *mut u64 as *mut _),
                std::mem::size_of::<u64>() as u32,
                Some(&mut bytes_returned),
                None,
            )?;
        }

        Ok(reply_data)
    }

    fn detect_cpu_vendor(&mut self) -> Result<()> {
        // Try Intel first
        self.vendor = CpuVendor::Intel;
        if self.get_data_from_driver(MSR_INTEL_RAPL_POWER_UNIT).is_ok() {
            return Ok(());
        }

        // Try AMD
        self.vendor = CpuVendor::Amd;
        if self.get_data_from_driver(MSR_AMD_RAPL_POWER_UNIT).is_ok() {
            return Ok(());
        }

        Err(Error::from_thread())
    }

    fn get_energy_units(&mut self) -> Result<()> {
        let msr = match self.vendor {
            CpuVendor::Intel => MSR_INTEL_RAPL_POWER_UNIT,
            CpuVendor::Amd => MSR_AMD_RAPL_POWER_UNIT,
        };

        let reply_data = self.get_data_from_driver(msr)?;

        // Time Units
        const TIME_MASK: u64 = 0xF0000;
        let time_val = reply_data & TIME_MASK;
        self.time_unit = 1.0 / 2.0_f64.powi((time_val >> 16) as i32);

        // Energy Units
        const ENERGY_MASK: u64 = 0x1F00;
        let energy_val = reply_data & ENERGY_MASK;
        self.energy_unit = 1.0 / 2.0_f64.powi((energy_val >> 8) as i32);

        // Power Units
        const POWER_MASK: u64 = 0xF;
        let power_val = reply_data & POWER_MASK;
        self.power_unit = 1.0 / 2.0_f64.powi(power_val as i32);

        // Reject degenerate units from a driver that returned unexpected MSR data.
        // Non-finite or non-positive units would cascade into NaN/infinity power.
        if !self.energy_unit.is_finite() || self.energy_unit <= 0.0 {
            return Err(Error::from_hresult(windows::core::HRESULT(
                0x80004005u32 as i32, // E_FAIL
            )));
        }

        Ok(())
    }

    fn check_supported_platform(&mut self) -> Result<()> {
        match self.vendor {
            CpuVendor::Intel => {
                // Intel: use only PKG energy
                if let Ok(reply_data) = self.get_data_from_driver(MSR_INTEL_PKG_ENERGY_STATUS)
                    && reply_data != 0
                {
                    self.pkg = true;
                }
            }
            CpuVendor::Amd => {
                // AMD only supports PKG, no PSYS
                if let Ok(reply_data) = self.get_data_from_driver(MSR_AMD_PKG_ENERGY_STATUS)
                    && reply_data != 0
                {
                    self.pkg = true;
                }
            }
        }

        Ok(())
    }

    pub fn get_rapl_energy(&self) -> Result<f64> {
        match self.vendor {
            CpuVendor::Intel => self.get_intel_energy(),
            CpuVendor::Amd => self.get_amd_energy(),
        }
    }

    fn get_intel_energy(&self) -> Result<f64> {
        if self.pkg {
            let reply_data = self.get_data_from_driver(MSR_INTEL_PKG_ENERGY_STATUS)?;
            let raw_pkg_energy = (reply_data & 0xFFFFFFFF) as u32;
            let pkg_energy = raw_pkg_energy as f64 * self.energy_unit;

            return Ok(pkg_energy);
        }

        Ok(0.0)
    }

    fn get_amd_energy(&self) -> Result<f64> {
        if self.pkg {
            let reply_data = self.get_data_from_driver(MSR_AMD_PKG_ENERGY_STATUS)?;
            let raw_pkg_energy = (reply_data & 0xFFFFFFFF) as u32;
            let pkg_energy = raw_pkg_energy as f64 * self.energy_unit;
            return Ok(pkg_energy);
        }

        Ok(0.0)
    }

    /// Calculate power since last time
    pub fn calculate_power(&mut self) -> io::Result<f64> {
        let current_energy: f64 = self.get_rapl_energy()?;
        let current_read_at = Instant::now();

        let Some(last_read_at) = self.last_read_at else {
            self.last_energy = current_energy;
            self.last_read_at = Some(current_read_at);
            return Ok(0.0);
        };

        let energy_delta = if current_energy >= self.last_energy {
            current_energy - self.last_energy
        } else {
            // PKG has wrapped. Since it's a 32-bit counter (0xFFFFFFFF),
            // we calculate the remaining energy before wrap and add current energy.
            let max_energy = 0xFFFFFFFFu64 as f64 * self.energy_unit;
            current_energy - self.last_energy + max_energy
        };
        let elapsed_secs = current_read_at.duration_since(last_read_at).as_secs_f64();

        self.last_energy = current_energy;
        self.last_read_at = Some(current_read_at);

        if elapsed_secs <= 0.0 || !elapsed_secs.is_finite() {
            return Ok(0.0);
        }

        Ok(energy_delta / elapsed_secs)
    }
}

impl Drop for RaplDriver {
    fn drop(&mut self) {
        unsafe {
            let _ = CloseHandle(self.handle);
        }
    }
}