#############################################################
#### Prometheus Windows Exporter Installer               ####
#### For Windows Server 2016+ and Windows 10/11          ####
####                                                     ####
#### Author: Phil Connor                                 ####
#### Contact: contact@mylinux.work                       ####
#### License: MIT                                        ####
#### Version: 1.0                                        ####
####                                                     ####
#### Usage: .\install-windows-exporter.ps1 [OPTIONS]     ####
#############################################################

param(
    [string]$ServerType = "standard",        # standard, iis, sql, ad, rdp, all
    [string]$ExtraCollectors = "",           # additional collectors (comma-separated)
    [string]$RemoveCollectors = "",          # collectors to remove from the profile
    [int]$Port = 9182,
    [string]$InstallDir = "C:\Program Files\windows_exporter",
    [string]$PrometheusIP = "",              # restrict firewall to this IP
    [switch]$NoFirewall,
    [switch]$NoDefenderExclusion,
    [switch]$Update,
    [switch]$Uninstall,
    [switch]$DryRun,
    [switch]$Help
)

$ErrorActionPreference = "Stop"

# ============================================================================
# CONFIGURATION
# ============================================================================

$ServiceName = "windows_exporter"
$ServiceDisplayName = "Prometheus Windows Exporter"
$FirewallRuleName = "Prometheus Windows Exporter (TCP-In)"
$GitHubApiUrl = "https://api.github.com/repos/prometheus-community/windows_exporter/releases/latest"
$TextFileDir = Join-Path $InstallDir "textfile_inputs"
$TargetExe = Join-Path $InstallDir "windows_exporter.exe"

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

function Write-ColorOutput {
    param([string]$Message, [string]$Color = "Green")
    $timestamp = Get-Date -Format "yyyy-MM-dd HH:mm:ss"
    Write-Host "[$timestamp] $Message" -ForegroundColor $Color
}

function Show-Help {
    Write-Host @"

Prometheus Windows Exporter Installer v1.0
https://mylinux.work

USAGE:
    .\install-windows-exporter.ps1 [OPTIONS]

OPTIONS:
    -ServerType <type>       Server profile: standard, iis, sql, ad, rdp, all (default: standard)
    -ExtraCollectors <list>  Additional collectors, comma-separated (e.g. "service,hyperv")
    -RemoveCollectors <list> Remove collectors from the profile (e.g. "tcp,time")
    -Port <port>             Metrics listen port (default: 9182)
    -InstallDir <path>       Installation directory (default: C:\Program Files\windows_exporter)
    -PrometheusIP <ip>       Restrict firewall rule to this source IP
    -NoFirewall              Skip firewall rule creation
    -NoDefenderExclusion     Skip Windows Defender exclusion
    -Update                  Update existing installation to latest version
    -Uninstall               Remove windows_exporter completely
    -DryRun                  Show what would be done without making changes
    -Help                    Show this help

SERVER PROFILES:
    standard    [defaults],process,tcp,time
    iis         [defaults],process,tcp,time,iis
    sql         [defaults],process,tcp,time,mssql
    ad          [defaults],process,tcp,time,ad,dns
    rdp         [defaults],process,tcp,time,remote_fx
    all         [defaults],process,tcp,time,iis,mssql,ad,dns,remote_fx

    [defaults] = cpu,cs,logical_disk,os,system,net,cache,logon,memory

EXAMPLES:
    # Standard install
    .\install-windows-exporter.ps1

    # IIS server with restricted firewall
    .\install-windows-exporter.ps1 -ServerType iis -PrometheusIP "10.0.0.5"

    # SQL server with extra service collector
    .\install-windows-exporter.ps1 -ServerType sql -ExtraCollectors "service"

    # Update to latest version
    .\install-windows-exporter.ps1 -Update

    # Uninstall
    .\install-windows-exporter.ps1 -Uninstall
"@
}

# ============================================================================
# ADMINISTRATOR CHECK
# ============================================================================

function Test-Administrator {
    $currentUser = [Security.Principal.WindowsIdentity]::GetCurrent()
    $principal = New-Object Security.Principal.WindowsPrincipal($currentUser)
    return $principal.IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)
}

# ============================================================================
# SERVER PROFILE BUILDER
# ============================================================================

function Get-CollectorString {
    $defaultCollectors = @("cpu", "cs", "logical_disk", "os", "system", "net", "cache", "logon", "memory")
    $profileCollectors = @("process", "tcp", "time")

    switch ($ServerType.ToLower()) {
        "iis"      { $profileCollectors += @("iis") }
        "sql"      { $profileCollectors += @("mssql") }
        "ad"       { $profileCollectors += @("ad", "dns") }
        "rdp"      { $profileCollectors += @("remote_fx") }
        "all"      { $profileCollectors += @("iis", "mssql", "ad", "dns", "remote_fx") }
        "standard" { }
        default {
            Write-ColorOutput "WARNING: Unknown server type '$ServerType', using 'standard'" "Yellow"
        }
    }

    $allCollectors = $defaultCollectors + $profileCollectors

    # Add extra collectors
    if ($ExtraCollectors) {
        $extras = $ExtraCollectors -split "," | ForEach-Object { $_.Trim() } | Where-Object { $_ -ne "" }
        foreach ($col in $extras) {
            if ($allCollectors -notcontains $col) {
                $allCollectors += $col
            }
        }
    }

    # Remove unwanted collectors
    if ($RemoveCollectors) {
        $removals = $RemoveCollectors -split "," | ForEach-Object { $_.Trim() } | Where-Object { $_ -ne "" }
        $allCollectors = $allCollectors | Where-Object { $removals -notcontains $_ }
    }

    return ($allCollectors | Select-Object -Unique) -join ","
}

# ============================================================================
# GITHUB RELEASE FUNCTIONS
# ============================================================================

function Get-LatestRelease {
    Write-ColorOutput "Querying GitHub for latest release..."
    try {
        $release = Invoke-RestMethod -Uri $GitHubApiUrl -UseBasicParsing
        $asset = $release.assets | Where-Object { $_.name -match "amd64" -and $_.name -match "\.exe$" } | Select-Object -First 1

        if (-not $asset) {
            Write-ColorOutput "ERROR: Could not find amd64 .exe asset in latest release" "Red"
            exit 1
        }

        return @{
            Version     = $release.tag_name
            DownloadUrl = $asset.browser_download_url
            AssetName   = $asset.name
        }
    } catch {
        Write-ColorOutput "ERROR: Failed to query GitHub API: $($_.Exception.Message)" "Red"
        exit 1
    }
}

# ============================================================================
# SERVICE MANAGEMENT
# ============================================================================

function Stop-ExistingService {
    $service = Get-Service -Name $ServiceName -ErrorAction SilentlyContinue
    if ($service) {
        Write-ColorOutput "Stopping existing service..."
        Stop-Service -Name $ServiceName -Force -ErrorAction SilentlyContinue
        Start-Sleep -Seconds 3
    }
}

function Remove-ExistingService {
    $service = Get-Service -Name $ServiceName -ErrorAction SilentlyContinue
    if ($service) {
        Stop-ExistingService
        Write-ColorOutput "Deleting existing service..."
        sc.exe delete $ServiceName | Out-Null
        Start-Sleep -Seconds 2
    }
}

function New-ExporterService {
    param([string]$CollectorString)

    $serviceArgs = @(
        "--log.file=eventlog",
        "--web.listen-address=0.0.0.0:$Port",
        "--collector.textfile.directories=`"$TextFileDir`"",
        "--collectors.enabled=$CollectorString"
    )

    $binaryPath = "`"$TargetExe`" " + ($serviceArgs -join " ")

    Write-ColorOutput "Creating service with collectors: $CollectorString" "Cyan"
    if ($DryRun) {
        Write-ColorOutput "[DRY RUN] Would create service: $binaryPath" "Yellow"
        return
    }

    New-Service -Name $ServiceName `
        -BinaryPathName $binaryPath `
        -DisplayName $ServiceDisplayName `
        -StartupType Automatic `
        -Description "Prometheus metrics exporter for Windows (managed by install-windows-exporter.ps1)" | Out-Null

    Write-ColorOutput "Service created successfully"
}

# ============================================================================
# FIREWALL MANAGEMENT
# ============================================================================

function Add-FirewallRule {
    if ($NoFirewall) {
        Write-ColorOutput "Skipping firewall rule (-NoFirewall specified)" "Yellow"
        return
    }

    # Remove existing rule if present
    $existing = Get-NetFirewallRule -DisplayName $FirewallRuleName -ErrorAction SilentlyContinue
    if ($existing) {
        Remove-NetFirewallRule -DisplayName $FirewallRuleName -ErrorAction SilentlyContinue
        Write-ColorOutput "Removed existing firewall rule"
    }

    if ($DryRun) {
        $scope = if ($PrometheusIP) { $PrometheusIP } else { "Any" }
        Write-ColorOutput "[DRY RUN] Would create firewall rule: port $Port, source: $scope" "Yellow"
        return
    }

    $params = @{
        DisplayName = $FirewallRuleName
        Direction   = "Inbound"
        Protocol    = "TCP"
        LocalPort   = $Port
        Action      = "Allow"
        Profile     = "Any"
        Description = "Allow Prometheus to scrape windows_exporter metrics on port $Port"
    }

    if ($PrometheusIP) {
        $params.RemoteAddress = $PrometheusIP
        Write-ColorOutput "Creating firewall rule: port $Port restricted to $PrometheusIP"
    } else {
        Write-ColorOutput "Creating firewall rule: port $Port open to all sources"
    }

    New-NetFirewallRule @params | Out-Null
    Write-ColorOutput "Firewall rule created"
}

function Remove-FirewallRule {
    $existing = Get-NetFirewallRule -DisplayName $FirewallRuleName -ErrorAction SilentlyContinue
    if ($existing) {
        Remove-NetFirewallRule -DisplayName $FirewallRuleName
        Write-ColorOutput "Firewall rule removed"
    } else {
        Write-ColorOutput "No firewall rule found to remove" "Yellow"
    }
}

# ============================================================================
# DEFENDER EXCLUSION
# ============================================================================

function Add-DefenderExclusion {
    if ($NoDefenderExclusion) {
        Write-ColorOutput "Skipping Defender exclusion (-NoDefenderExclusion specified)" "Yellow"
        return
    }

    if ($DryRun) {
        Write-ColorOutput "[DRY RUN] Would add Defender exclusion for: $InstallDir" "Yellow"
        return
    }

    try {
        Add-MpPreference -ExclusionPath $InstallDir -ErrorAction Stop
        Write-ColorOutput "Windows Defender exclusion added for: $InstallDir"
    } catch {
        Write-ColorOutput "WARNING: Could not add Defender exclusion: $($_.Exception.Message)" "Yellow"
    }
}

function Remove-DefenderExclusion {
    try {
        Remove-MpPreference -ExclusionPath $InstallDir -ErrorAction Stop
        Write-ColorOutput "Windows Defender exclusion removed for: $InstallDir"
    } catch {
        Write-ColorOutput "WARNING: Could not remove Defender exclusion: $($_.Exception.Message)" "Yellow"
    }
}

# ============================================================================
# VERIFICATION
# ============================================================================

function Test-MetricsEndpoint {
    Write-ColorOutput "Waiting for service to start..."
    Start-Sleep -Seconds 5

    $metricsUrl = "http://localhost:$Port/metrics"
    try {
        $response = Invoke-WebRequest -Uri $metricsUrl -UseBasicParsing -TimeoutSec 10
        if ($response.StatusCode -eq 200) {
            Write-ColorOutput "Metrics endpoint responding at $metricsUrl" "Green"
            return $true
        }
    } catch {
        Write-ColorOutput "WARNING: Metrics endpoint not responding at $metricsUrl" "Yellow"
        Write-ColorOutput "The service may still be starting. Check: Get-Service $ServiceName" "Yellow"
        return $false
    }
    return $false
}

# ============================================================================
# UNINSTALL MODE
# ============================================================================

function Invoke-Uninstall {
    Write-ColorOutput "=== Uninstalling Windows Exporter ===" "Cyan"

    $service = Get-Service -Name $ServiceName -ErrorAction SilentlyContinue
    if (-not $service -and -not (Test-Path $InstallDir)) {
        Write-ColorOutput "Windows Exporter does not appear to be installed" "Yellow"
        exit 0
    }

    Write-Host ""
    Write-Host "This will remove:" -ForegroundColor Yellow
    Write-Host "  - Service: $ServiceName"
    Write-Host "  - Directory: $InstallDir"
    Write-Host "  - Firewall rule: $FirewallRuleName"
    Write-Host "  - Defender exclusion: $InstallDir"
    Write-Host ""

    if (-not $DryRun) {
        $confirm = Read-Host "Are you sure? (y/N)"
        if ($confirm -ne "y" -and $confirm -ne "Y") {
            Write-ColorOutput "Uninstall cancelled" "Yellow"
            exit 0
        }
    }

    if ($DryRun) {
        Write-ColorOutput "[DRY RUN] Would uninstall windows_exporter" "Yellow"
        return
    }

    # Stop and remove service
    Remove-ExistingService

    # Remove firewall rule
    Remove-FirewallRule

    # Remove Defender exclusion
    Remove-DefenderExclusion

    # Remove install directory
    if (Test-Path $InstallDir) {
        Remove-Item -Path $InstallDir -Recurse -Force
        Write-ColorOutput "Install directory removed: $InstallDir"
    }

    Write-Host ""
    Write-ColorOutput "=== Windows Exporter uninstalled successfully ===" "Green"
}

# ============================================================================
# UPDATE MODE
# ============================================================================

function Invoke-Update {
    Write-ColorOutput "=== Updating Windows Exporter ===" "Cyan"

    $service = Get-Service -Name $ServiceName -ErrorAction SilentlyContinue
    if (-not $service) {
        Write-ColorOutput "ERROR: Service '$ServiceName' not found. Use install mode instead." "Red"
        exit 1
    }

    $release = Get-LatestRelease
    Write-ColorOutput "Latest version: $($release.Version)"

    # Check installed version if binary exists
    if (Test-Path $TargetExe) {
        try {
            $installedVersion = (Get-Item $TargetExe).VersionInfo.ProductVersion
            if ($installedVersion) {
                Write-ColorOutput "Installed version: $installedVersion"
                $latestClean = $release.Version -replace '^v', ''
                if ($installedVersion -eq $latestClean) {
                    Write-ColorOutput "Already running the latest version. No update needed." "Green"
                    exit 0
                }
            }
        } catch {
            Write-ColorOutput "Could not determine installed version, proceeding with update" "Yellow"
        }
    }

    if ($DryRun) {
        Write-ColorOutput "[DRY RUN] Would update to $($release.Version)" "Yellow"
        return
    }

    # Download new binary
    $tempExe = Join-Path $env:TEMP "windows_exporter_update.exe"
    Write-ColorOutput "Downloading $($release.AssetName)..."
    Invoke-WebRequest -Uri $release.DownloadUrl -OutFile $tempExe -UseBasicParsing

    # Stop service, replace binary, start service
    Stop-ExistingService
    Copy-Item -Path $tempExe -Destination $TargetExe -Force
    Remove-Item $tempExe -Force -ErrorAction SilentlyContinue
    Write-ColorOutput "Binary updated"

    Start-Service -Name $ServiceName
    $verified = Test-MetricsEndpoint

    Write-Host ""
    Write-ColorOutput "=== Update Summary ===" "Cyan"
    Write-ColorOutput "Version:      $($release.Version)"
    Write-ColorOutput "Status:       $(if ($verified) { 'Verified' } else { 'Started (verification pending)' })"
    Write-ColorOutput "Metrics URL:  http://localhost:$Port/metrics"
}

# ============================================================================
# INSTALL MODE
# ============================================================================

function Invoke-Install {
    Write-ColorOutput "=== Installing Windows Exporter ===" "Cyan"
    Write-ColorOutput "Server type:  $ServerType"
    Write-ColorOutput "Port:         $Port"
    Write-ColorOutput "Install dir:  $InstallDir"

    # --- Create directories ---
    if (-not $DryRun) {
        if (-not (Test-Path $InstallDir)) {
            New-Item -Path $InstallDir -ItemType Directory -Force | Out-Null
            Write-ColorOutput "Created directory: $InstallDir"
        }
        if (-not (Test-Path $TextFileDir)) {
            New-Item -Path $TextFileDir -ItemType Directory -Force | Out-Null
            Write-ColorOutput "Created directory: $TextFileDir"
        }
    } else {
        Write-ColorOutput "[DRY RUN] Would create: $InstallDir" "Yellow"
        Write-ColorOutput "[DRY RUN] Would create: $TextFileDir" "Yellow"
    }

    # --- Defender exclusion ---
    Add-DefenderExclusion

    # --- Get latest release ---
    $release = Get-LatestRelease
    Write-ColorOutput "Latest version: $($release.Version)"

    # --- Check if already installed with same version ---
    if (Test-Path $TargetExe) {
        try {
            $installedVersion = (Get-Item $TargetExe).VersionInfo.ProductVersion
            $latestClean = $release.Version -replace '^v', ''
            if ($installedVersion -and $installedVersion -eq $latestClean) {
                Write-ColorOutput "Version $installedVersion already installed, skipping download" "Yellow"
                $skipDownload = $true
            }
        } catch { }
    }

    # --- Download binary ---
    $tempExe = Join-Path $env:TEMP "windows_exporter_install.exe"
    if (-not $skipDownload) {
        Write-ColorOutput "Downloading $($release.AssetName)..."
        if (-not $DryRun) {
            Invoke-WebRequest -Uri $release.DownloadUrl -OutFile $tempExe -UseBasicParsing
            Write-ColorOutput "Download complete"
        } else {
            Write-ColorOutput "[DRY RUN] Would download: $($release.DownloadUrl)" "Yellow"
        }
    }

    # --- Stop and remove existing service ---
    Remove-ExistingService

    # --- Copy binary ---
    if (-not $skipDownload -and -not $DryRun) {
        Copy-Item -Path $tempExe -Destination $TargetExe -Force
        Remove-Item $tempExe -Force -ErrorAction SilentlyContinue
        Write-ColorOutput "Binary installed to: $TargetExe"
    }

    # --- Build collector string ---
    $collectorString = Get-CollectorString

    # --- Create service ---
    New-ExporterService -CollectorString $collectorString

    # --- Firewall rule ---
    Add-FirewallRule

    # --- Start service ---
    if (-not $DryRun) {
        Write-ColorOutput "Starting service..."
        Start-Service -Name $ServiceName
        $verified = Test-MetricsEndpoint
    }

    # --- Summary ---
    Write-Host ""
    Write-ColorOutput "=== Installation Summary ===" "Cyan"
    Write-ColorOutput "Version:      $($release.Version)"
    Write-ColorOutput "Binary:       $TargetExe"
    Write-ColorOutput "Port:         $Port"
    Write-ColorOutput "Collectors:   $collectorString"
    Write-ColorOutput "Textfile dir: $TextFileDir"
    Write-ColorOutput "Firewall:     $(if ($NoFirewall) { 'Skipped' } elseif ($PrometheusIP) { "Restricted to $PrometheusIP" } else { 'Open' })"
    Write-ColorOutput "Defender:     $(if ($NoDefenderExclusion) { 'Skipped' } else { 'Exclusion added' })"

    if (-not $DryRun) {
        Write-ColorOutput "Status:       $(if ($verified) { 'Verified - metrics endpoint responding' } else { 'Started (verification pending)' })"
        Write-ColorOutput "Verify URL:   http://localhost:$Port/metrics" "Cyan"
    }

    Write-Host ""
    Write-ColorOutput "=== Installation complete ===" "Green"
}

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

# Show help
if ($Help) {
    Show-Help
    exit 0
}

# Check administrator privileges
if (-not (Test-Administrator)) {
    Write-ColorOutput "ERROR: This script must be run as Administrator" "Red"
    Write-Host "Right-click PowerShell and select 'Run as administrator'" -ForegroundColor Yellow
    exit 1
}

if ($DryRun) {
    Write-ColorOutput "=== DRY RUN MODE - No changes will be made ===" "Yellow"
    Write-Host ""
}

# Route to the appropriate mode
if ($Uninstall) {
    Invoke-Uninstall
} elseif ($Update) {
    Invoke-Update
} else {
    Invoke-Install
}