<#
.SYNOPSIS
    Disables Windows power-saving features (including hibernate on low battery),
    applies a synthetic CPU load, and logs battery drain over time.

.DESCRIPTION
    1. Re-launches itself elevated if not running as Administrator (powercfg changes require admin).
    2. Disables sleep/hibernate/monitor timeouts, USB selective suspend, PCI Express link power
       management, and the "hibernate/shutdown on critical/low battery" actions, on both AC and DC.
    3. Sets the active power plan to High performance (or Ultimate Performance if available).
    4. Starts a background synthetic load across all logical processors targeting ~30% CPU usage
       (configurable via -TargetLoadPercent).
    5. Samples battery percentage/status every -SampleIntervalSeconds and appends to a CSV log
       until -DurationMinutes elapses (or indefinitely with -DurationMinutes 0, stop with Ctrl+C).
    6. Prints a summary (start/end charge, elapsed time, average drain rate %/hour).

.PARAMETER DurationMinutes
    How long to run the test. 0 = run until Ctrl+C is pressed.

.PARAMETER SampleIntervalSeconds
    How often to sample/log the battery status.

.PARAMETER TargetLoadPercent
    Approximate overall CPU load to generate (0-100). Default 30.

.PARAMETER LogPath
    Path to the CSV log file. Defaults to .\BatteryDrainLog_<timestamp>.csv next to the script.

.PARAMETER SkipPowerSettingsChanges
    Skip modifying power settings; only run the load + logging.

.EXAMPLE
    .\Test-BatteryDrain.ps1 -DurationMinutes 120 -TargetLoadPercent 30

.NOTES
    Restoring defaults: the script does not automatically revert power settings it changes.
    Run `powercfg /restoredefaultschemes` (drastic) or re-enable individual settings manually
    if you want the previous behavior back.
#>

[CmdletBinding()]
param(
    [int]$DurationMinutes = 60,
    [int]$SampleIntervalSeconds = 30,
    [ValidateRange(0, 100)]
    [int]$TargetLoadPercent = 30,
    [string]$LogPath = $(Join-Path -Path $PSScriptRoot -ChildPath ("BatteryDrainLog_{0:yyyyMMdd_HHmmss}.csv" -f (Get-Date))),
    [string]$StateFilePath = $(Join-Path -Path $PSScriptRoot -ChildPath 'PowerSettingsBackup.json'),
    [switch]$SkipPowerSettingsChanges
)

$ErrorActionPreference = 'Stop'

#region Elevation check
function Test-IsAdmin {
    $identity = [Security.Principal.WindowsIdentity]::GetCurrent()
    $principal = New-Object Security.Principal.WindowsPrincipal($identity)
    return $principal.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
}

if (-not (Test-IsAdmin)) {
    Write-Host "Not running elevated - relaunching as Administrator..." -ForegroundColor Yellow
    $argList = @(
        '-NoProfile', '-ExecutionPolicy', 'Bypass', '-File', "`"$PSCommandPath`"",
        '-DurationMinutes', $DurationMinutes,
        '-SampleIntervalSeconds', $SampleIntervalSeconds,
        '-TargetLoadPercent', $TargetLoadPercent,
        '-LogPath', "`"$LogPath`"",
        '-StateFilePath', "`"$StateFilePath`""
    )
    if ($SkipPowerSettingsChanges) { $argList += '-SkipPowerSettingsChanges' }
    Start-Process -FilePath 'powershell.exe' -ArgumentList $argList -Verb RunAs
    exit
}
#endregion

#region Power settings
function Save-PowerSettingsState {
    param([string]$Path)

    # Original active scheme GUID (strip the trailing description in parentheses).
    $activeLine = (powercfg /getactivescheme) -join ''
    $activeGuid = if ($activeLine -match '([0-9a-fA-F-]{36})') { $Matches[1] } else { $null }

    # Hibernate is considered enabled if hiberfil.sys exists / powercfg /a doesn't list it as unavailable due to being off.
    $hibernateEnabled = (powercfg /a) -join "`n" -notmatch 'Hibernation has not been enabled'

    $state = [PSCustomObject]@{
        OriginalSchemeGuid = $activeGuid
        HibernateEnabled   = [bool]$hibernateEnabled
        SavedAt            = (Get-Date).ToString('o')
    }
    $state | ConvertTo-Json | Out-File -FilePath $Path -Encoding UTF8
    Write-Host "Saved original power settings to $Path (used by Restore-PowerSettings.ps1)." -ForegroundColor DarkGray
}

function Set-PowerSettingsForBatteryTest {
    Write-Host "Configuring power plan to prevent sleep/hibernate/screen-off and low-battery actions..." -ForegroundColor Cyan

    Save-PowerSettingsState -Path $StateFilePath

    # Try to use Ultimate/High performance plan so Windows itself doesn't throttle or sleep.
    $schemes = powercfg /list
    $ultimateGuid = ($schemes | Select-String -Pattern 'Ultimate Performance' | ForEach-Object {
        ($_ -match '([0-9a-fA-F-]{36})') | Out-Null; $Matches[1]
    }) | Select-Object -First 1

    if (-not $ultimateGuid) {
        try {
            $dup = powercfg -duplicatescheme e9a42b02-d5df-448d-aa00-03f14749eb61 2>$null
            $ultimateGuid = ($dup | Select-String -Pattern '([0-9a-fA-F-]{36})' | ForEach-Object { $Matches[1] }) | Select-Object -First 1
        } catch { }
    }

    if ($ultimateGuid) {
        powercfg /setactive $ultimateGuid | Out-Null
        Write-Host "Active plan set to Ultimate Performance." -ForegroundColor Green
    } else {
        # SCHEME_MIN = High performance
        powercfg /setactive SCHEME_MIN | Out-Null
        Write-Host "Active plan set to High performance." -ForegroundColor Green
    }

    # Disable timeouts for sleep, hibernate, monitor, disk - both AC (ac) and DC (battery, dc).
    $targets = @('monitor-timeout-ac', 'monitor-timeout-dc',
                 'disk-timeout-ac', 'disk-timeout-dc',
                 'standby-timeout-ac', 'standby-timeout-dc',
                 'hibernate-timeout-ac', 'hibernate-timeout-dc')
    foreach ($t in $targets) {
        powercfg /change $t 0 | Out-Null
    }

    # Disable "Allow hybrid sleep" and standby/hibernate at the GUID level (belt & suspenders).
    # SUB_SLEEP settings
    $SUB_SLEEP = '238c9fa8-0aad-41ed-83f4-97be242c8f20'
    $STANDBYIDLE = '29f6c1db-86da-48c5-9fdb-f2b67b1f44da'
    $HIBERNATEIDLE = '9d7815a6-7ee4-497e-8888-515a05f02364'
    foreach ($mode in @('/setacvalueindex', '/setdcvalueindex')) {
        powercfg $mode SCHEME_CURRENT $SUB_SLEEP $STANDBYIDLE 0 | Out-Null
        powercfg $mode SCHEME_CURRENT $SUB_SLEEP $HIBERNATEIDLE 0 | Out-Null
    }

    # Disable low/critical battery actions (hibernate/shutdown) - the actual ask: no hibernate on low battery.
    $SUB_BATTERY = 'e73a048d-bf27-4f12-9731-8b2076e8891f'
    $BATACTIONCRIT = '637ea02f-bbcb-4015-8e2c-a1c793b9d808' # Critical battery action
    $BATACTIONLOW  = 'd8742dcb-3e6a-4b3c-b3fe-374623cdcf06' # Low battery action
    foreach ($mode in @('/setacvalueindex', '/setdcvalueindex')) {
        # 0 = Do Nothing, 1 = Sleep, 2 = Hibernate, 3 = Shut Down
        powercfg $mode SCHEME_CURRENT $SUB_BATTERY $BATACTIONCRIT 0 | Out-Null
        powercfg $mode SCHEME_CURRENT $SUB_BATTERY $BATACTIONLOW 0 | Out-Null
    }

    # Turn off notifications tied to low battery flyout as well (0 = off).
    $BATFLAGSCRIT = '5dbb7c9f-38e9-40d2-9749-4f8a0e9f640f'
    $BATFLAGSLOW  = 'bcded951-187b-4d05-bccc-f7e51960c258'
    foreach ($mode in @('/setacvalueindex', '/setdcvalueindex')) {
        powercfg $mode SCHEME_CURRENT $SUB_BATTERY $BATFLAGSCRIT 0 2>$null | Out-Null
        powercfg $mode SCHEME_CURRENT $SUB_BATTERY $BATFLAGSLOW 0 2>$null | Out-Null
    }

    # Explicitly disable OS-level hibernate (removes hiberfil.sys, guarantees no hibernate trigger).
    powercfg /hibernate off | Out-Null

    # Disable USB selective suspend and PCI Express link state power management (both drain-reducers, but they
    # can throttle performance under sustained load, which we don't want while stress-testing).
    $SUB_USB = '2a737441-1930-4402-8d77-b2bebba308a3'
    $USBSUSPEND = '48e6b7a6-50f5-4782-a5d4-53bb8f07e226'
    $SUB_PCIE = '501a4d13-42af-4429-9fd1-a8218c268e20'
    $PCIE_ASPM = 'ee12f906-d277-404b-b6da-e5fa1a576df5'
    foreach ($mode in @('/setacvalueindex', '/setdcvalueindex')) {
        powercfg $mode SCHEME_CURRENT $SUB_USB $USBSUSPEND 0 2>$null | Out-Null
        powercfg $mode SCHEME_CURRENT $SUB_PCIE $PCIE_ASPM 0 2>$null | Out-Null
    }

    # Disable adaptive brightness / dimming on battery.
    $SUB_VIDEO = '7516b95f-f776-4464-8c53-06167f40cc99'
    $ADAPTBRIGHT = 'fbd9aa66-9553-4097-ba44-ed6e9d65eab8'
    $DIMDISPLAY = '17aaa29b-8b43-4b94-aafe-35f64daaf1ee'
    foreach ($mode in @('/setacvalueindex', '/setdcvalueindex')) {
        powercfg $mode SCHEME_CURRENT $SUB_VIDEO $ADAPTBRIGHT 0 2>$null | Out-Null
        powercfg $mode SCHEME_CURRENT $SUB_VIDEO $DIMDISPLAY 0 2>$null | Out-Null
    }

    powercfg /setactive SCHEME_CURRENT | Out-Null

    # Prevent Windows from turning off the display/sleeping via the SendMessage/SetThreadExecutionState API too,
    # as a redundant safety net for the duration of this script (handled later via a background job).

    Write-Host "Power settings applied." -ForegroundColor Green
}
#endregion

#region Keep system awake (SetThreadExecutionState) for the life of the script
Add-Type -Namespace Native -Name Power -MemberDefinition @'
[DllImport("kernel32.dll", CharSet = CharSet.Auto, SetLastError = true)]
public static extern uint SetThreadExecutionState(uint esFlags);
'@

$ES_CONTINUOUS = [Convert]::ToUInt32('80000000', 16)
$ES_SYSTEM_REQUIRED = [Convert]::ToUInt32('00000001', 16)
$ES_DISPLAY_REQUIRED = [Convert]::ToUInt32('00000002', 16)

function Set-StayAwake {
    [Native.Power]::SetThreadExecutionState([uint32]($ES_CONTINUOUS -bor $ES_SYSTEM_REQUIRED -bor $ES_DISPLAY_REQUIRED)) | Out-Null
}

function Clear-StayAwake {
    [Native.Power]::SetThreadExecutionState($ES_CONTINUOUS) | Out-Null
}
#endregion

#region Synthetic CPU load (~TargetLoadPercent across all logical cores)
function Start-CpuLoad {
    param([int]$Percent, [int]$Cores)

    $jobs = @()
    for ($i = 0; $i -lt $Cores; $i++) {
        $jobs += Start-Job -ScriptBlock {
            param($pct)
            $busyMs = [double]$pct
            $idleMs = 100.0 - $busyMs
            $sw = [System.Diagnostics.Stopwatch]::new()
            while ($true) {
                $sw.Restart()
                while ($sw.Elapsed.TotalMilliseconds -lt $busyMs) {
                    # Busy-spin to consume CPU
                    [Math]::Sqrt(12345.6789) | Out-Null
                }
                if ($idleMs -gt 0) {
                    Start-Sleep -Milliseconds $idleMs
                }
            }
        } -ArgumentList $Percent
    }
    return $jobs
}

function Stop-CpuLoad {
    param($Jobs)
    if ($Jobs) {
        $Jobs | Stop-Job -PassThru | Remove-Job -Force
    }
}
#endregion

#region Battery sampling
function Get-BatteryInfo {
    $battery = Get-CimInstance -ClassName Win32_Battery -ErrorAction SilentlyContinue | Select-Object -First 1
    if (-not $battery) {
        return [PSCustomObject]@{
            Timestamp        = Get-Date
            PercentRemaining = $null
            Status           = 'No battery detected'
            Charging         = $null
            EstimatedRunTime = $null
        }
    }

    $chargingStatus = switch ($battery.BatteryStatus) {
        1 { 'Discharging' }
        2 { 'AC Power (Charged)' }
        3 { 'Fully Charged' }
        4 { 'Low' }
        5 { 'Critical' }
        6 { 'Charging' }
        7 { 'Charging High' }
        8 { 'Charging Low' }
        9 { 'Charging Critical' }
        10 { 'Undefined' }
        11 { 'Partially Charged' }
        default { "Unknown ($($battery.BatteryStatus))" }
    }

    [PSCustomObject]@{
        Timestamp        = Get-Date
        PercentRemaining = $battery.EstimatedChargeRemaining
        Status           = $chargingStatus
        Charging         = ($battery.BatteryStatus -in 2, 6, 7, 8, 9)
        EstimatedRunTime  = $battery.EstimatedRunTime
    }
}
#endregion

#region Main
try {
    # Validate/prepare the log path up front so a bad path fails fast, before power settings or
    # the CPU load are touched. Falls back to the script folder if the requested location is invalid.
    $logDir = Split-Path -Path $LogPath -Parent
    if ($logDir -and -not (Test-Path -Path $logDir)) {
        try {
            New-Item -ItemType Directory -Path $logDir -Force | Out-Null
        } catch {
            $fallbackPath = Join-Path -Path $PSScriptRoot -ChildPath (Split-Path -Path $LogPath -Leaf)
            Write-Warning "Could not create log directory '$logDir' ($($_.Exception.Message)). Falling back to '$fallbackPath'."
            $LogPath = $fallbackPath
        }
    }

    if (-not $SkipPowerSettingsChanges) {
        Set-PowerSettingsForBatteryTest
    } else {
        Write-Host "Skipping power settings changes (per -SkipPowerSettingsChanges)." -ForegroundColor Yellow
    }

    Set-StayAwake

    $cores = [Environment]::ProcessorCount
    Write-Host "Starting synthetic CPU load targeting ~$TargetLoadPercent% across $cores logical processors..." -ForegroundColor Cyan
    $loadJobs = Start-CpuLoad -Percent $TargetLoadPercent -Cores $cores

    "Timestamp,PercentRemaining,Status,Charging,EstimatedRunTimeMinutes" | Out-File -FilePath $LogPath -Encoding UTF8

    Write-Host "Logging battery status every $SampleIntervalSeconds s to $LogPath" -ForegroundColor Cyan
    if ($DurationMinutes -gt 0) {
        Write-Host "Test will run for $DurationMinutes minute(s). Press Ctrl+C to stop early." -ForegroundColor Cyan
    } else {
        Write-Host "Test will run indefinitely. Press Ctrl+C to stop." -ForegroundColor Cyan
    }

    $startTime = Get-Date
    $endTime = if ($DurationMinutes -gt 0) { $startTime.AddMinutes($DurationMinutes) } else { [DateTime]::MaxValue }
    $firstSample = $null
    $lastSample = $null

    while ((Get-Date) -lt $endTime) {
        $info = Get-BatteryInfo
        if (-not $firstSample) { $firstSample = $info }
        $lastSample = $info

        $line = "{0:o},{1},{2},{3},{4}" -f $info.Timestamp, $info.PercentRemaining, $info.Status, $info.Charging, $info.EstimatedRunTime
        Add-Content -Path $LogPath -Value $line

        Write-Host ("[{0:HH:mm:ss}] Battery: {1}%  Status: {2}" -f $info.Timestamp, $info.PercentRemaining, $info.Status)

        Start-Sleep -Seconds $SampleIntervalSeconds
    }
}
finally {
    Write-Host "Stopping CPU load and restoring system idle behavior..." -ForegroundColor Cyan
    Stop-CpuLoad -Jobs $loadJobs
    Clear-StayAwake

    if ($firstSample -and $lastSample -and $firstSample.PercentRemaining -ne $null -and $lastSample.PercentRemaining -ne $null) {
        $elapsedHours = ($lastSample.Timestamp - $firstSample.Timestamp).TotalHours
        $drop = $firstSample.PercentRemaining - $lastSample.PercentRemaining
        $rate = if ($elapsedHours -gt 0) { $drop / $elapsedHours } else { 0 }

        Write-Host ""
        Write-Host "===== Battery Drain Test Summary =====" -ForegroundColor Yellow
        Write-Host ("Start:  {0:yyyy-MM-dd HH:mm:ss}  {1}%" -f $firstSample.Timestamp, $firstSample.PercentRemaining)
        Write-Host ("End:    {0:yyyy-MM-dd HH:mm:ss}  {1}%" -f $lastSample.Timestamp, $lastSample.PercentRemaining)
        Write-Host ("Elapsed: {0:N2} hour(s)" -f $elapsedHours)
        Write-Host ("Total drop: {0} percentage points" -f $drop)
        Write-Host ("Average drain rate: {0:N2} %/hour" -f $rate)
        Write-Host ("Log file: {0}" -f $LogPath)
        Write-Host "=======================================" -ForegroundColor Yellow
    }

    if (-not $SkipPowerSettingsChanges) {
        Write-Host "Run .\Restore-PowerSettings.ps1 to revert the power/hibernate settings changed by this script." -ForegroundColor Cyan
    }
}
#endregion