#Requires -Version 5.1
#Requires -RunAsAdministrator
<#
.SYNOPSIS
    Updates PowerShell modules in the repository for WinPE integration.

.DESCRIPTION
    Downloads or updates PowerShell modules from PSGallery to the modules/ directory.
    These modules are then committed to the repository and included in WinPE builds.

.PARAMETER ModuleName
    Name of a specific module to update. If not specified, updates all modules in config.json.

.PARAMETER Force
    Forces re-download even if module already exists.

.EXAMPLE
    .\Update-Modules.ps1
    Updates all modules listed in config.json

.EXAMPLE
    .\Update-Modules.ps1 -ModuleName "PSWindowsUpdate" -Force
    Forces update of specific module
#>

[CmdletBinding()]
param(
    [Parameter(Mandatory = $false)]
    [string]$ModuleName,

    [Parameter(Mandatory = $false)]
    [switch]$Force
)

$ErrorActionPreference = 'Stop'

$ConfigPath = Join-Path $PSScriptRoot "config.json"
$ModulesRoot = Join-Path $PSScriptRoot "modules"

function Write-Status {
    param(
        [string]$Message,
        [string]$Type = "Info"
    )

    $color = switch ($Type) {
        "Success" { "Green" }
        "Warning" { "Yellow" }
        "Error" { "Red" }
        "Info" { "Cyan" }
        default { "White" }
    }

    $prefix = switch ($Type) {
        "Success" { "[+]" }
        "Warning" { "[!]" }
        "Error" { "[X]" }
        "Info" { "[i]" }
        default { "[-]" }
    }

    Write-Host "$prefix " -ForegroundColor $color -NoNewline
    Write-Host $Message
}

function Get-ConfiguredModules {
    if (-not (Test-Path $ConfigPath)) {
        throw "Configuration file not found: $ConfigPath"
    }

    try {
        $config = Get-Content $ConfigPath -Raw | ConvertFrom-Json
        return $config.powershellModules
    }
    catch {
        throw "Failed to read configuration: $_"
    }
}

function Initialize-ModulesDirectory {
    if (-not (Test-Path $ModulesRoot)) {
        Write-Status "Creating modules directory: $ModulesRoot" "Info"
        New-Item -ItemType Directory -Path $ModulesRoot | Out-Null
    }
}

function Update-PSGalleryModule {
    param(
        [Parameter(Mandatory = $true)]
        [string]$Name
    )

    Write-Host ""
    Write-Host "=" * 70 -ForegroundColor Cyan
    Write-Status "Processing module: $Name" "Info"
    Write-Host "=" * 70 -ForegroundColor Cyan

    $modulePath = Join-Path $ModulesRoot $Name

    if ((Test-Path $modulePath) -and -not $Force) {
        Write-Status "Module already exists: $modulePath" "Warning"
        Write-Status "Use -Force to re-download" "Info"
        return
    }

    if (Test-Path $modulePath) {
        Write-Status "Removing existing module..." "Warning"
        Remove-Item $modulePath -Recurse -Force
    }

    Write-Status "Searching PSGallery for module: $Name" "Info"

    try {
        $moduleInfo = Find-Module -Name $Name -Repository PSGallery -ErrorAction Stop

        Write-Status "Found: $Name v$($moduleInfo.Version)" "Success"
        Write-Status "Description: $($moduleInfo.Description)" "Info"

        if ($moduleInfo.Dependencies -and $moduleInfo.Dependencies.Count -gt 0) {
            Write-Status "Dependencies detected:" "Warning"
            foreach ($dep in $moduleInfo.Dependencies) {
                Write-Host "  - $($dep.Name)" -ForegroundColor Yellow
            }
            Write-Status "Dependencies will be downloaded automatically" "Info"
        }

        Write-Status "Downloading module to repository..." "Info"

        Save-Module -Name $Name -Path $ModulesRoot -Repository PSGallery -Force

        if (Test-Path $modulePath) {
            $savedVersion = (Get-ChildItem $modulePath -Directory | Select-Object -First 1).Name
            Write-Status "Successfully saved: $Name v$savedVersion" "Success"

            $moduleSizeMB = (Get-ChildItem $modulePath -Recurse | Measure-Object -Property Length -Sum).Sum / 1MB
            Write-Status "Module size: $([math]::Round($moduleSizeMB, 2)) MB" "Info"
        }
        else {
            throw "Module save completed but directory not found"
        }
    }
    catch {
        Write-Status "Failed to update module: $Name" "Error"
        Write-Status "Error: $($_.Exception.Message)" "Error"
        throw
    }
}

try {
    Write-Host ""
    Write-Host "================================================" -ForegroundColor Cyan
    Write-Host "  PowerShell Module Repository Updater" -ForegroundColor Yellow
    Write-Host "================================================" -ForegroundColor Cyan
    Write-Host ""

    Initialize-ModulesDirectory

    if ($ModuleName) {
        Write-Status "Single module update requested: $ModuleName" "Info"
        $modulesToUpdate = @($ModuleName)
    }
    else {
        Write-Status "Reading configuration: $ConfigPath" "Info"
        $modulesToUpdate = Get-ConfiguredModules

        if ($null -eq $modulesToUpdate -or $modulesToUpdate.Count -eq 0) {
            Write-Status "No modules configured in config.json" "Warning"
            Write-Status "Add modules to 'powershellModules' array in config.json" "Info"
            exit 0
        }

        Write-Status "Found $($modulesToUpdate.Count) module(s) in configuration" "Success"
    }

    Write-Host ""

    foreach ($module in $modulesToUpdate) {
        Update-PSGalleryModule -Name $module
    }

    Write-Host ""
    Write-Host "================================================" -ForegroundColor Green
    Write-Status "Module update completed successfully" "Success"
    Write-Host "================================================" -ForegroundColor Green
    Write-Host ""
    Write-Status "Modules saved to: $ModulesRoot" "Info"
    Write-Status "These modules will be included in the next WinPE build" "Info"
    Write-Host ""

    Write-Host "Next steps:" -ForegroundColor Yellow
    Write-Host "  1. Review downloaded modules in: $ModulesRoot" -ForegroundColor Cyan
    Write-Host "  2. Commit changes to git if desired" -ForegroundColor Cyan
    Write-Host "  3. Run Build-PE.ps1 to create updated WinPE image" -ForegroundColor Cyan
    Write-Host ""
}
catch {
    Write-Host ""
    Write-Status "FATAL ERROR: $($_.Exception.Message)" "Error"
    Write-Host ""
    exit 1
}