<#
.SYNOPSIS
    Microsoft Deployment Toolkit Prometheus Metrics Exporter
.DESCRIPTION
    Prometheus exporter for MDT - active deployments, completed/failed task
    sequences, deployment duration, OS image count, task sequence count,
    application count, driver count, driver package count, boot image count,
    deployment share size. Exports metrics as Prometheus-compatible text format
    for the windows_exporter textfile collector.
.PARAMETER TextfilePath
    Full path to the .prom output file (default: C:\metrics\mdt.prom)
.PARAMETER DeploymentShare
    UNC or local path to the MDT deployment share
.PARAMETER MDTDatabase
    Path to MDT monitoring database file (optional)
.PARAMETER InstallScheduledTask
    Switch to create a scheduled task for periodic collection
.PARAMETER TaskIntervalMinutes
    Interval in minutes for the scheduled task (default: 5)
.NOTES
    Author:  Phil Connor
    Contact: contact@mylinux.work
    Website: https://mylinux.work
    License: MIT
    Version: 1.0.0

    Metrics Exported:
      Status:
        - mdt_up
        - mdt_info{version}

      Deployments:
        - mdt_deployment_active
        - mdt_deployment_completed_total
        - mdt_deployment_failed_total
        - mdt_deployment_duration_seconds_min
        - mdt_deployment_duration_seconds_max
        - mdt_deployment_duration_seconds_avg

      Inventory:
        - mdt_os_image_count
        - mdt_task_sequence_count
        - mdt_application_count
        - mdt_driver_count
        - mdt_driver_package_count

      Boot Images:
        - mdt_boot_image_count
        - mdt_boot_image_size_bytes

      Capacity:
        - mdt_deployment_share_size_bytes

      Exporter:
        - mdt_exporter_duration_seconds
        - mdt_exporter_last_run_timestamp
#>

param(
    [string]$TextfilePath = 'C:\metrics\mdt.prom',

    [string]$DeploymentShare = '\\localhost\DeploymentShare$',

    [string]$MDTDatabase = '',

    [switch]$InstallScheduledTask,

    [int]$TaskIntervalMinutes = 5
)

# ============================================================================
# SCHEDULED TASK INSTALLATION
# ============================================================================

if ($InstallScheduledTask) {
    $taskName = "MdtMetricsExporter"
    $existingTask = Get-ScheduledTask -TaskName $taskName -ErrorAction SilentlyContinue

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

        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 MDT 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"
    }
    return
}

$ErrorActionPreference = 'SilentlyContinue'
$VERSION = '1.0.0'

# ============================================================================
# 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 Resolve-SharePath {
    param([string]$SharePath)
    if (Test-Path $SharePath) {
        return (Resolve-Path $SharePath).Path
    }
    return $SharePath
}

# ============================================================================
# DEPLOYMENT SHARE INVENTORY METRICS
# ============================================================================

function Get-InventoryMetrics {
    param([string]$ShareRoot)

    $sb = [System.Text.StringBuilder]::new()

    # --- mdt_os_image_count ---
    [void]$sb.AppendLine('# HELP mdt_os_image_count Number of OS images in the deployment share')
    [void]$sb.AppendLine('# TYPE mdt_os_image_count gauge')
    try {
        $osPath = Join-Path $ShareRoot 'Operating Systems'
        if (Test-Path $osPath) {
            $wimFiles = Get-ChildItem -Path $osPath -Recurse -File -Filter '*.wim' -ErrorAction Stop
            $osCount = ($wimFiles | Measure-Object).Count
        } else {
            $osCount = 0
        }
        [void]$sb.AppendLine("mdt_os_image_count $osCount")
    } catch {
        [void]$sb.AppendLine("mdt_os_image_count 0")
    }
    [void]$sb.AppendLine('')

    # --- mdt_task_sequence_count ---
    [void]$sb.AppendLine('# HELP mdt_task_sequence_count Number of task sequences defined')
    [void]$sb.AppendLine('# TYPE mdt_task_sequence_count gauge')
    try {
        $tsPath = Join-Path $ShareRoot 'Control'
        if (Test-Path $tsPath) {
            $tsFolders = Get-ChildItem -Path $tsPath -Directory -ErrorAction Stop |
                Where-Object { Test-Path (Join-Path $_.FullName 'ts.xml') }
            $tsCount = ($tsFolders | Measure-Object).Count
        } else {
            $tsCount = 0
        }
        [void]$sb.AppendLine("mdt_task_sequence_count $tsCount")
    } catch {
        [void]$sb.AppendLine("mdt_task_sequence_count 0")
    }
    [void]$sb.AppendLine('')

    # --- mdt_application_count ---
    [void]$sb.AppendLine('# HELP mdt_application_count Number of applications in the deployment share')
    [void]$sb.AppendLine('# TYPE mdt_application_count gauge')
    try {
        $appPath = Join-Path $ShareRoot 'Applications'
        if (Test-Path $appPath) {
            $appFolders = Get-ChildItem -Path $appPath -Directory -ErrorAction Stop
            $appCount = ($appFolders | Measure-Object).Count
        } else {
            $appCount = 0
        }
        [void]$sb.AppendLine("mdt_application_count $appCount")
    } catch {
        [void]$sb.AppendLine("mdt_application_count 0")
    }
    [void]$sb.AppendLine('')

    # --- mdt_driver_count ---
    [void]$sb.AppendLine('# HELP mdt_driver_count Number of individual drivers imported')
    [void]$sb.AppendLine('# TYPE mdt_driver_count gauge')
    try {
        $driverPath = Join-Path $ShareRoot 'Out-of-Box Drivers'
        if (Test-Path $driverPath) {
            $infFiles = Get-ChildItem -Path $driverPath -Recurse -File -Filter '*.inf' -ErrorAction Stop
            $driverCount = ($infFiles | Measure-Object).Count
        } else {
            $driverCount = 0
        }
        [void]$sb.AppendLine("mdt_driver_count $driverCount")
    } catch {
        [void]$sb.AppendLine("mdt_driver_count 0")
    }
    [void]$sb.AppendLine('')

    # --- mdt_driver_package_count ---
    [void]$sb.AppendLine('# HELP mdt_driver_package_count Number of driver package folders')
    [void]$sb.AppendLine('# TYPE mdt_driver_package_count gauge')
    try {
        $driverPath = Join-Path $ShareRoot 'Out-of-Box Drivers'
        if (Test-Path $driverPath) {
            $driverPkgs = Get-ChildItem -Path $driverPath -Directory -ErrorAction Stop
            $driverPkgCount = ($driverPkgs | Measure-Object).Count
        } else {
            $driverPkgCount = 0
        }
        [void]$sb.AppendLine("mdt_driver_package_count $driverPkgCount")
    } catch {
        [void]$sb.AppendLine("mdt_driver_package_count 0")
    }
    [void]$sb.AppendLine('')

    $sb.ToString()
}

# ============================================================================
# BOOT IMAGE METRICS
# ============================================================================

function Get-BootImageMetrics {
    param([string]$ShareRoot)

    $sb = [System.Text.StringBuilder]::new()

    # --- mdt_boot_image_count ---
    [void]$sb.AppendLine('# HELP mdt_boot_image_count Number of boot images (WIM and ISO)')
    [void]$sb.AppendLine('# TYPE mdt_boot_image_count gauge')
    $bootCount = 0
    $bootSizeTotal = 0
    try {
        $bootPath = Join-Path $ShareRoot 'Boot'
        if (Test-Path $bootPath) {
            $bootFiles = Get-ChildItem -Path $bootPath -File -ErrorAction Stop |
                Where-Object { $_.Extension -in '.wim', '.iso' }
            $bootCount = ($bootFiles | Measure-Object).Count
            $bootSizeTotal = ($bootFiles | Measure-Object -Property Length -Sum).Sum
            if (-not $bootSizeTotal) { $bootSizeTotal = 0 }
        }
    } catch {
        $bootCount = 0
        $bootSizeTotal = 0
    }
    [void]$sb.AppendLine("mdt_boot_image_count $bootCount")
    [void]$sb.AppendLine('')

    # --- mdt_boot_image_size_bytes ---
    [void]$sb.AppendLine('# HELP mdt_boot_image_size_bytes Total size of boot images in bytes')
    [void]$sb.AppendLine('# TYPE mdt_boot_image_size_bytes gauge')
    [void]$sb.AppendLine("mdt_boot_image_size_bytes $bootSizeTotal")
    [void]$sb.AppendLine('')

    $sb.ToString()
}

# ============================================================================
# DEPLOYMENT SHARE SIZE
# ============================================================================

function Get-ShareSizeMetrics {
    param([string]$ShareRoot)

    $sb = [System.Text.StringBuilder]::new()

    # --- mdt_deployment_share_size_bytes ---
    [void]$sb.AppendLine('# HELP mdt_deployment_share_size_bytes Total deployment share size in bytes')
    [void]$sb.AppendLine('# TYPE mdt_deployment_share_size_bytes gauge')
    try {
        if (Test-Path $ShareRoot) {
            $totalSize = (Get-ChildItem -Path $ShareRoot -Recurse -File -ErrorAction Stop |
                Measure-Object -Property Length -Sum).Sum
            if (-not $totalSize) { $totalSize = 0 }
        } else {
            $totalSize = 0
        }
        [void]$sb.AppendLine("mdt_deployment_share_size_bytes $totalSize")
    } catch {
        [void]$sb.AppendLine("mdt_deployment_share_size_bytes 0")
    }
    [void]$sb.AppendLine('')

    $sb.ToString()
}

# ============================================================================
# MDT MONITORING / DEPLOYMENT STATUS METRICS
# ============================================================================

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

    $activeCount = 0
    $completedCount = 0
    $failedCount = 0
    $durations = @()
    $monitorAvailable = $false

    # Attempt to load MDT PowerShell module and query monitoring data
    try {
        $mdtModule = Get-Module -ListAvailable -Name MicrosoftDeploymentToolkit -ErrorAction Stop
        if ($mdtModule) {
            Import-Module MicrosoftDeploymentToolkit -ErrorAction Stop
            $monitorData = Get-MDTMonitorData -Path "DS001:" -ErrorAction Stop
            $monitorAvailable = $true

            foreach ($entry in $monitorData) {
                switch ($entry.DeploymentStatus) {
                    1 { $activeCount++
                        if ($entry.StartTime) {
                            $elapsed = ((Get-Date) - [datetime]$entry.StartTime).TotalSeconds
                            $durations += $elapsed
                        }
                    }
                    2 { $completedCount++ }
                    3 { $failedCount++ }
                }
            }
        }
    } catch {
        $monitorAvailable = $false
    }

    # If MDT module not available, try reading monitoring database directly
    if (-not $monitorAvailable -and $MDTDatabase -and (Test-Path $MDTDatabase)) {
        try {
            $connectionString = "Data Source=$MDTDatabase;Version=3;"
            $connection = New-Object System.Data.SQLite.SQLiteConnection($connectionString)
            $connection.Open()

            $cmd = $connection.CreateCommand()
            $cmd.CommandText = "SELECT DeploymentStatus, StartTime FROM MonitorData"
            $reader = $cmd.ExecuteReader()

            while ($reader.Read()) {
                $status = $reader['DeploymentStatus']
                switch ($status) {
                    1 {
                        $activeCount++
                        $startStr = $reader['StartTime']
                        if ($startStr) {
                            $elapsed = ((Get-Date) - [datetime]$startStr).TotalSeconds
                            $durations += $elapsed
                        }
                    }
                    2 { $completedCount++ }
                    3 { $failedCount++ }
                }
            }

            $reader.Close()
            $connection.Close()
            $monitorAvailable = $true
        } catch {
            $monitorAvailable = $false
        }
    }

    # --- mdt_deployment_active ---
    [void]$sb.AppendLine('# HELP mdt_deployment_active Number of currently active MDT deployments')
    [void]$sb.AppendLine('# TYPE mdt_deployment_active gauge')
    [void]$sb.AppendLine("mdt_deployment_active $activeCount")
    [void]$sb.AppendLine('')

    # --- mdt_deployment_completed_total ---
    [void]$sb.AppendLine('# HELP mdt_deployment_completed_total Total completed MDT deployments')
    [void]$sb.AppendLine('# TYPE mdt_deployment_completed_total counter')
    [void]$sb.AppendLine("mdt_deployment_completed_total $completedCount")
    [void]$sb.AppendLine('')

    # --- mdt_deployment_failed_total ---
    [void]$sb.AppendLine('# HELP mdt_deployment_failed_total Total failed MDT deployments')
    [void]$sb.AppendLine('# TYPE mdt_deployment_failed_total counter')
    [void]$sb.AppendLine("mdt_deployment_failed_total $failedCount")
    [void]$sb.AppendLine('')

    # --- deployment duration stats ---
    if ($durations.Count -gt 0) {
        $minDuration = Format-MetricValue ($durations | Measure-Object -Minimum).Minimum
        $maxDuration = Format-MetricValue ($durations | Measure-Object -Maximum).Maximum
        $avgDuration = Format-MetricValue ($durations | Measure-Object -Average).Average
    } else {
        $minDuration = 0
        $maxDuration = 0
        $avgDuration = 0
    }

    [void]$sb.AppendLine('# HELP mdt_deployment_duration_seconds_min Shortest active deployment duration')
    [void]$sb.AppendLine('# TYPE mdt_deployment_duration_seconds_min gauge')
    [void]$sb.AppendLine("mdt_deployment_duration_seconds_min $minDuration")
    [void]$sb.AppendLine('')

    [void]$sb.AppendLine('# HELP mdt_deployment_duration_seconds_max Longest active deployment duration')
    [void]$sb.AppendLine('# TYPE mdt_deployment_duration_seconds_max gauge')
    [void]$sb.AppendLine("mdt_deployment_duration_seconds_max $maxDuration")
    [void]$sb.AppendLine('')

    [void]$sb.AppendLine('# HELP mdt_deployment_duration_seconds_avg Average active deployment duration')
    [void]$sb.AppendLine('# TYPE mdt_deployment_duration_seconds_avg gauge')
    [void]$sb.AppendLine("mdt_deployment_duration_seconds_avg $avgDuration")
    [void]$sb.AppendLine('')

    $sb.ToString()
}

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

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

    $shareRoot = Resolve-SharePath $DeploymentShare

    # --- mdt_up ---
    [void]$sb.AppendLine('# HELP mdt_up MDT exporter status (1=share accessible, 0=error)')
    [void]$sb.AppendLine('# TYPE mdt_up gauge')
    if (Test-Path $shareRoot) {
        [void]$sb.AppendLine("mdt_up 1")
    } else {
        [void]$sb.AppendLine("mdt_up 0")
    }
    [void]$sb.AppendLine('')

    # --- mdt_info ---
    [void]$sb.AppendLine('# HELP mdt_info MDT exporter version information')
    [void]$sb.AppendLine('# TYPE mdt_info gauge')
    [void]$sb.AppendLine("mdt_info{version=`"$VERSION`"} 1")
    [void]$sb.AppendLine('')

    # Deployment status metrics
    [void]$sb.Append((Get-DeploymentMetrics))

    # Inventory metrics
    [void]$sb.Append((Get-InventoryMetrics -ShareRoot $shareRoot))

    # Boot image metrics
    [void]$sb.Append((Get-BootImageMetrics -ShareRoot $shareRoot))

    # Share size metrics
    [void]$sb.Append((Get-ShareSizeMetrics -ShareRoot $shareRoot))

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

    [void]$sb.AppendLine('# HELP mdt_exporter_duration_seconds Time to generate all metrics')
    [void]$sb.AppendLine('# TYPE mdt_exporter_duration_seconds gauge')
    [void]$sb.AppendLine("mdt_exporter_duration_seconds $duration")
    [void]$sb.AppendLine('')

    [void]$sb.AppendLine('# HELP mdt_exporter_last_run_timestamp Unix timestamp of last successful run')
    [void]$sb.AppendLine('# TYPE mdt_exporter_last_run_timestamp gauge')
    [void]$sb.AppendLine("mdt_exporter_last_run_timestamp $timestamp")
    [void]$sb.AppendLine('')

    $sb.ToString()
}

# ============================================================================
# MAIN EXECUTION - TEXTFILE OUTPUT
# ============================================================================

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

$tempFile = Join-Path $outputDir ".mdt_metrics.$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 $TextfilePath -Force
    Write-Host "Metrics written to $TextfilePath ($lineCount lines)" -ForegroundColor Green
}
catch {
    Remove-Item -Path $tempFile -Force -ErrorAction SilentlyContinue
    Write-Error "Failed to generate metrics: $_"
    exit 1
}