<#
.SYNOPSIS
    Windows BITS Transfer Prometheus Metrics Exporter
.DESCRIPTION
    Prometheus exporter for Windows BITS transfer metrics - active transfers,
    bytes transferred, job states, bandwidth throttling, completed jobs per
    hour, and error counts. Exports metrics as Prometheus-compatible text format.
.PARAMETER Mode
    Output mode: 'stdout' (default), 'textfile', or 'http'
.PARAMETER Port
    HTTP port for http mode (default: 9518)
.PARAMETER TextfileDir
    Directory for textfile collector output (default: C:\ProgramData\node_exporter)
.PARAMETER InstallScheduledTask
    Switch to create a scheduled task for auto-start on system boot
.PARAMETER TaskIntervalMinutes
    Interval in minutes for the scheduled task (default: 2)
.NOTES
    Author:  Phil Connor
    Contact: contact@mylinux.work
    Website: https://mylinux.work
    License: MIT
    Version: 1.0

    Metrics Exported:
      Core Status:
        - bits_up
        - bits_exporter_info{version}

      Jobs:
        - bits_jobs_total
        - bits_jobs_transferring
        - bits_jobs_queued
        - bits_jobs_suspended
        - bits_jobs_error
        - bits_jobs_completed

      Transfer:
        - bits_bytes_transferred_total
        - bits_bytes_remaining_total

      Bandwidth:
        - bits_bandwidth_limit_bytes

      Performance:
        - bits_completed_jobs_per_hour

      Exporter:
        - bits_exporter_duration_seconds
        - bits_exporter_last_run_timestamp
#>

param(
    [ValidateSet('stdout', 'textfile', 'http')]
    [string]$Mode = 'stdout',

    [int]$Port = 9518,

    [string]$TextfileDir = 'C:\ProgramData\node_exporter',

    [switch]$InstallScheduledTask,

    [int]$TaskIntervalMinutes = 2
)

# Create a scheduled task to run this script every $TaskIntervalMinutes minutes
if ($InstallScheduledTask) {
    $taskName = "BitsTransferExporter"
    $existingTask = Get-ScheduledTask -TaskName $taskName -ErrorAction SilentlyContinue

    if (-not $existingTask) {
        $taskAction = New-ScheduledTaskAction -Execute "powershell.exe" -Argument "-NoProfile -ExecutionPolicy Bypass -File `"$($MyInvocation.MyCommand.Path)`" -Mode textfile"

        if (-not $TaskIntervalMinutes -or $TaskIntervalMinutes -le 0) {
            throw "TaskIntervalMinutes must be a positive integer"
        }

        $taskTrigger = New-ScheduledTaskTrigger -Once -At (Get-Date).AddMinutes(1) -RepetitionInterval (New-TimeSpan -Minutes $TaskIntervalMinutes) -RepetitionDuration (New-TimeSpan -Days 365)
        $taskPrincipal = New-ScheduledTaskPrincipal -UserId "SYSTEM" -LogonType ServiceAccount -RunLevel Highest

        try {
            Write-Host "Creating scheduled task: $taskName"
            Register-ScheduledTask -TaskName $taskName -Action $taskAction -Trigger $taskTrigger -Principal $taskPrincipal -Description "Exports BITS transfer metrics for Prometheus every $TaskIntervalMinutes minutes"

            $createdTask = Get-ScheduledTask -TaskName $taskName -ErrorAction SilentlyContinue
            if (-not $createdTask) {
                throw "Failed to verify scheduled task creation"
            }
            Write-Host "Successfully created scheduled task: $taskName" -ForegroundColor Green
        } catch {
            Write-Error "Failed to create scheduled task: $($_.Exception.Message)"
            throw
        }
    } else {
        Write-Host "Scheduled task '$taskName' already exists, skipping creation"
    }
}

$ErrorActionPreference = 'SilentlyContinue'

# ============================================================================
# HELPER FUNCTIONS
# ============================================================================

function Get-UnixTimestamp {
    [int][double]::Parse((Get-Date -UFormat '%s'))
}

function Format-MetricValue {
    param([double]$Value, [int]$Decimals = 2)
    [math]::Round($Value, $Decimals)
}

function Get-SystemUptimeHours {
    try {
        $os = Get-CimInstance -ClassName Win32_OperatingSystem -ErrorAction Stop
        $uptime = (Get-Date) - $os.LastBootUpTime
        if ($uptime.TotalHours -lt 0.01) { return 0.01 }
        return $uptime.TotalHours
    } catch {
        return 1
    }
}

# ============================================================================
# BITS SERVICE STATUS
# ============================================================================

function Get-BitsServiceStatus {
    try {
        $service = Get-Service -Name BITS -ErrorAction Stop
        if ($service.Status -eq 'Running') {
            return $true
        }
        return $false
    } catch {
        return $false
    }
}

# ============================================================================
# BITS BANDWIDTH POLICY
# ============================================================================

function Get-BitsBandwidthLimit {
    try {
        $regPath = 'HKLM:\SOFTWARE\Policies\Microsoft\Windows\BITS'
        if (Test-Path $regPath) {
            $maxBandwidth = Get-ItemProperty -Path $regPath -Name 'MaxBandwidthServed' -ErrorAction SilentlyContinue
            if ($maxBandwidth -and $maxBandwidth.MaxBandwidthServed) {
                return [long]$maxBandwidth.MaxBandwidthServed
            }

            $maxTransferRate = Get-ItemProperty -Path $regPath -Name 'MaxTransferRateOffSchedule' -ErrorAction SilentlyContinue
            if ($maxTransferRate -and $maxTransferRate.MaxTransferRateOffSchedule) {
                return [long]$maxTransferRate.MaxTransferRateOffSchedule
            }
        }

        $regPathThrottle = 'HKLM:\SOFTWARE\Policies\Microsoft\Windows\BITS\Throttle'
        if (Test-Path $regPathThrottle) {
            $bandwidthLimit = Get-ItemProperty -Path $regPathThrottle -Name 'BandwidthLimit' -ErrorAction SilentlyContinue
            if ($bandwidthLimit -and $bandwidthLimit.BandwidthLimit) {
                return [long]$bandwidthLimit.BandwidthLimit
            }
        }

        return 0
    } catch {
        return 0
    }
}

# ============================================================================
# BITS JOB METRICS
# ============================================================================

function Get-BitsJobMetrics {
    $result = @{
        Total        = 0
        Transferring = 0
        Queued       = 0
        Suspended    = 0
        Error        = 0
        Completed    = 0
        BytesTransferred = [long]0
        BytesRemaining   = [long]0
    }

    try {
        Import-Module BitsTransfer -ErrorAction Stop
    } catch {
        Write-Warning "BitsTransfer module not available: $_"
        return $result
    }

    # Get all BITS jobs for all users (requires elevation)
    $allJobs = @()
    try {
        $allJobs = @(Get-BitsTransfer -AllUsers -ErrorAction Stop)
    } catch {
        # Fall back to current user only
        try {
            $allJobs = @(Get-BitsTransfer -ErrorAction Stop)
        } catch {
            Write-Warning "Failed to retrieve BITS jobs: $_"
            return $result
        }
    }

    $result.Total = $allJobs.Count

    foreach ($job in $allJobs) {
        switch ($job.JobState) {
            'Transferring' {
                $result.Transferring++
            }
            'Connecting' {
                $result.Transferring++
            }
            'Queued' {
                $result.Queued++
            }
            'Suspended' {
                $result.Suspended++
            }
            'Error' {
                $result.Error++
            }
            'TransientError' {
                $result.Error++
            }
            'Transferred' {
                $result.Completed++
            }
            'Acknowledged' {
                $result.Completed++
            }
            'Cancelled' {
                # Cancelled jobs are not counted in active states
            }
            default {
                # Unknown states are ignored
            }
        }

        # Accumulate transfer bytes
        try {
            if ($job.BytesTransferred -ge 0) {
                $result.BytesTransferred += [long]$job.BytesTransferred
            }
        } catch {}

        try {
            if ($job.BytesTotal -ge 0 -and $job.BytesTransferred -ge 0) {
                $remaining = [long]$job.BytesTotal - [long]$job.BytesTransferred
                if ($remaining -gt 0) {
                    $result.BytesRemaining += $remaining
                }
            }
        } catch {}
    }

    return $result
}

# ============================================================================
# WMI BITS COMPLETED JOBS COUNT
# ============================================================================

function Get-BitsCompletedJobsFromEventLog {
    try {
        $filter = @{
            LogName   = 'Microsoft-Windows-Bits-Client/Operational'
            Id        = 4
            StartTime = (Get-Date).AddHours(-24)
        }
        $events = @(Get-WinEvent -FilterHashtable $filter -ErrorAction Stop)
        return $events.Count
    } catch {
        return 0
    }
}

# ============================================================================
# BITS METRICS COLLECTOR
# ============================================================================

function Get-BitsMetrics {
    $sb = [System.Text.StringBuilder]::new()

    # Check BITS service status
    $bitsUp = Get-BitsServiceStatus

    # --- bits_up ---
    [void]$sb.AppendLine('# HELP bits_up BITS service reachability (1=up, 0=down)')
    [void]$sb.AppendLine('# TYPE bits_up gauge')
    $upVal = if ($bitsUp) { 1 } else { 0 }
    [void]$sb.AppendLine("bits_up $upVal")
    [void]$sb.AppendLine('')

    # --- bits_exporter_info ---
    [void]$sb.AppendLine('# HELP bits_exporter_info Exporter version information')
    [void]$sb.AppendLine('# TYPE bits_exporter_info gauge')
    [void]$sb.AppendLine('bits_exporter_info{version="1.0"} 1')
    [void]$sb.AppendLine('')

    if (-not $bitsUp) {
        return $sb.ToString()
    }

    # Collect BITS job metrics
    $jobMetrics = Get-BitsJobMetrics

    # --- bits_jobs_total ---
    [void]$sb.AppendLine('# HELP bits_jobs_total Total BITS jobs')
    [void]$sb.AppendLine('# TYPE bits_jobs_total gauge')
    [void]$sb.AppendLine("bits_jobs_total $($jobMetrics.Total)")
    [void]$sb.AppendLine('')

    # --- bits_jobs_transferring ---
    [void]$sb.AppendLine('# HELP bits_jobs_transferring Actively transferring jobs')
    [void]$sb.AppendLine('# TYPE bits_jobs_transferring gauge')
    [void]$sb.AppendLine("bits_jobs_transferring $($jobMetrics.Transferring)")
    [void]$sb.AppendLine('')

    # --- bits_jobs_queued ---
    [void]$sb.AppendLine('# HELP bits_jobs_queued Queued jobs waiting to transfer')
    [void]$sb.AppendLine('# TYPE bits_jobs_queued gauge')
    [void]$sb.AppendLine("bits_jobs_queued $($jobMetrics.Queued)")
    [void]$sb.AppendLine('')

    # --- bits_jobs_suspended ---
    [void]$sb.AppendLine('# HELP bits_jobs_suspended Suspended jobs')
    [void]$sb.AppendLine('# TYPE bits_jobs_suspended gauge')
    [void]$sb.AppendLine("bits_jobs_suspended $($jobMetrics.Suspended)")
    [void]$sb.AppendLine('')

    # --- bits_jobs_error ---
    [void]$sb.AppendLine('# HELP bits_jobs_error Jobs in error state')
    [void]$sb.AppendLine('# TYPE bits_jobs_error gauge')
    [void]$sb.AppendLine("bits_jobs_error $($jobMetrics.Error)")
    [void]$sb.AppendLine('')

    # --- bits_jobs_completed ---
    [void]$sb.AppendLine('# HELP bits_jobs_completed Completed jobs since last reboot')
    [void]$sb.AppendLine('# TYPE bits_jobs_completed gauge')
    $completedFromLog = Get-BitsCompletedJobsFromEventLog
    $totalCompleted = $jobMetrics.Completed + $completedFromLog
    [void]$sb.AppendLine("bits_jobs_completed $totalCompleted")
    [void]$sb.AppendLine('')

    # --- bits_bytes_transferred_total ---
    [void]$sb.AppendLine('# HELP bits_bytes_transferred_total Total bytes transferred across all jobs')
    [void]$sb.AppendLine('# TYPE bits_bytes_transferred_total gauge')
    [void]$sb.AppendLine("bits_bytes_transferred_total $($jobMetrics.BytesTransferred)")
    [void]$sb.AppendLine('')

    # --- bits_bytes_remaining_total ---
    [void]$sb.AppendLine('# HELP bits_bytes_remaining_total Total bytes remaining across all jobs')
    [void]$sb.AppendLine('# TYPE bits_bytes_remaining_total gauge')
    [void]$sb.AppendLine("bits_bytes_remaining_total $($jobMetrics.BytesRemaining)")
    [void]$sb.AppendLine('')

    # --- bits_bandwidth_limit_bytes ---
    [void]$sb.AppendLine('# HELP bits_bandwidth_limit_bytes Configured BITS bandwidth limit in bytes per second')
    [void]$sb.AppendLine('# TYPE bits_bandwidth_limit_bytes gauge')
    $bandwidthLimit = Get-BitsBandwidthLimit
    [void]$sb.AppendLine("bits_bandwidth_limit_bytes $bandwidthLimit")
    [void]$sb.AppendLine('')

    # --- bits_completed_jobs_per_hour ---
    [void]$sb.AppendLine('# HELP bits_completed_jobs_per_hour Job completion rate per hour')
    [void]$sb.AppendLine('# TYPE bits_completed_jobs_per_hour gauge')
    $uptimeHours = Get-SystemUptimeHours
    $completedPerHour = Format-MetricValue ($totalCompleted / $uptimeHours)
    [void]$sb.AppendLine("bits_completed_jobs_per_hour $completedPerHour")
    [void]$sb.AppendLine('')

    $sb.ToString()
}

# ============================================================================
# COLLECT ALL METRICS
# ============================================================================

function Get-AllMetrics {
    $scriptStart = Get-Date
    $sb = [System.Text.StringBuilder]::new()

    # Collect BITS metrics
    [void]$sb.Append((Get-BitsMetrics))

    # Exporter runtime
    $scriptEnd = Get-Date
    $duration = Format-MetricValue ($scriptEnd - $scriptStart).TotalSeconds
    $timestamp = Get-UnixTimestamp

    [void]$sb.AppendLine('# HELP bits_exporter_duration_seconds Time to generate all metrics')
    [void]$sb.AppendLine('# TYPE bits_exporter_duration_seconds gauge')
    [void]$sb.AppendLine("bits_exporter_duration_seconds $duration")
    [void]$sb.AppendLine('')
    [void]$sb.AppendLine('# HELP bits_exporter_last_run_timestamp Unix timestamp of last successful run')
    [void]$sb.AppendLine('# TYPE bits_exporter_last_run_timestamp gauge')
    [void]$sb.AppendLine("bits_exporter_last_run_timestamp $timestamp")
    [void]$sb.AppendLine('')

    $sb.ToString()
}

# ============================================================================
# HTTP SERVER MODE
# ============================================================================

function Start-HttpServer {
    param([int]$ListenPort)

    $prefix = "http://+:$ListenPort/"
    $listener = [System.Net.HttpListener]::new()
    $listener.Prefixes.Add($prefix)

    try {
        $listener.Start()
        Write-Host "Starting BITS Transfer exporter on port $ListenPort..." -ForegroundColor Green
        Write-Host "Metrics available at http://localhost:$ListenPort/metrics"

        while ($listener.IsListening) {
            $context  = $listener.GetContext()
            $request  = $context.Request
            $response = $context.Response

            if ($request.Url.AbsolutePath -eq '/metrics') {
                $metrics = Get-AllMetrics
                $buffer  = [System.Text.Encoding]::UTF8.GetBytes($metrics)
                $response.ContentType = 'text/plain; version=0.0.4; charset=utf-8'
            }
            else {
                $html = @"
<!DOCTYPE html>
<html>
<head><title>BITS Transfer Exporter v1.0</title></head>
<body>
<h1>BITS Transfer Exporter v1.0</h1>
<p><a href="/metrics">Metrics</a></p>
<h2>Metrics</h2>
<ul>
<li>BITS service status</li>
<li>Job states (transferring, queued, suspended, error, completed)</li>
<li>Bytes transferred and remaining</li>
<li>Bandwidth throttling configuration</li>
<li>Completed jobs per hour rate</li>
</ul>
</body>
</html>
"@
                $buffer = [System.Text.Encoding]::UTF8.GetBytes($html)
                $response.ContentType = 'text/html; charset=utf-8'
            }

            $response.ContentLength64 = $buffer.Length
            $response.OutputStream.Write($buffer, 0, $buffer.Length)
            $response.OutputStream.Close()
        }
    }
    catch {
        Write-Error "HTTP server error: $_"
        Write-Error "If access denied, run: netsh http add urlacl url=http://+:$ListenPort/ user=Everyone"
    }
    finally {
        if ($listener.IsListening) {
            $listener.Stop()
        }
    }
}

# ============================================================================
# MAIN EXECUTION
# ============================================================================

switch ($Mode) {
    'http' {
        Start-HttpServer -ListenPort $Port
    }
    'textfile' {
        $OutputFile = Join-Path $TextfileDir 'bits_transfer.prom'

        $outputDir = Split-Path $OutputFile -Parent
        if (-not (Test-Path $outputDir)) {
            New-Item -Path $outputDir -ItemType Directory -Force | Out-Null
        }

        $tempFile = Join-Path $outputDir ".bits_transfer.$PID.tmp"

        try {
            $metrics = Get-AllMetrics
            $metrics | Out-File -FilePath $tempFile -Encoding utf8 -NoNewline

            $lineCount = ($metrics -split "`n").Count
            if ($lineCount -lt 10) {
                Remove-Item -Path $tempFile -Force -ErrorAction SilentlyContinue
                Write-Error "Metrics file too small ($lineCount lines), keeping previous"
                exit 1
            }

            Move-Item -Path $tempFile -Destination $OutputFile -Force
            Write-Host "Metrics written to $OutputFile ($lineCount lines)" -ForegroundColor Green
        }
        catch {
            Remove-Item -Path $tempFile -Force -ErrorAction SilentlyContinue
            Write-Error "Failed to generate metrics: $_"
            exit 1
        }
    }
    default {
        Get-AllMetrics | Write-Output
    }
}