#Requires -Version 5.1
<#
.SYNOPSIS
    Built-in Script Menu for WinPE Script Launcher
.DESCRIPTION
    Discovers and executes scripts from a .scripts folder. This is a simplified
    version designed for WinPE environments without UAC elevation requirements.
#>

[CmdletBinding()]
param(
    [string]$ScriptsPath = ".scripts"
)

$Script:SupportedExtensions = @('.ps1', '.bat', '.cmd')
$Script:ExitCode = 0

function Write-Header {
    param([string]$Title)

    $width = 70
    Write-Host ""
    Write-Host ("=" * $width) -ForegroundColor Cyan
    Write-Host "  $Title" -ForegroundColor Yellow
    Write-Host ("=" * $width) -ForegroundColor Cyan
    Write-Host ""
}

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

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

    $icon = switch ($Type) {
        'Success' { '[+]' }
        'Warning' { '[!]' }
        'Error' { '[X]' }
        default { '[i]' }
    }

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

function Write-Separator {
    Write-Host ("-" * 70) -ForegroundColor DarkGray
}

function Get-ScriptFiles {
    param([string]$Path)

    if (-not (Test-Path $Path)) {
        return $null
    }

    $scripts = Get-ChildItem -Path $Path -File -ErrorAction SilentlyContinue |
        Where-Object { $Script:SupportedExtensions -contains $_.Extension } |
        Sort-Object Name

    return $scripts
}

function Show-ScriptMenu {
    param([array]$Scripts)

    Write-Host "  Available Scripts:" -ForegroundColor Green
    Write-Host ""

    for ($i = 0; $i -lt $Scripts.Count; $i++) {
        $num = $i + 1
        $script = $Scripts[$i]
        $ext = $script.Extension

        $extColor = switch ($ext) {
            '.ps1' { 'Magenta' }
            '.bat' { 'Yellow' }
            '.cmd' { 'Cyan' }
            default { 'White' }
        }

        Write-Host "    " -NoNewline
        Write-Host ("[{0,2}]" -f $num) -ForegroundColor White -NoNewline
        Write-Host " " -NoNewline
        Write-Host $script.Name -ForegroundColor $extColor
    }

    Write-Host ""
    Write-Separator
    Write-Host ""
}

function Get-UserSelection {
    param([int]$MaxNumber)

    Write-Host "  Selection Options:" -ForegroundColor Yellow
    Write-Host "    - Enter number (e.g., '1')" -ForegroundColor Gray
    Write-Host "    - Enter 'all' to run all scripts sequentially" -ForegroundColor Gray
    Write-Host "    - Enter 'q' or 'quit' to exit" -ForegroundColor Gray
    Write-Host ""

    while ($true) {
        Write-Host "  Select script to run: " -ForegroundColor Cyan -NoNewline
        $input = Read-Host

        if ([string]::IsNullOrWhiteSpace($input)) {
            Write-InfoBox "Please enter a selection" "Warning"
            continue
        }

        $input = $input.Trim().ToLower()

        if ($input -eq 'q' -or $input -eq 'quit') {
            return $null
        }

        if ($input -eq 'all') {
            return 1..$MaxNumber
        }

        if ($input -match '^\d+$') {
            $num = [int]$input

            if ($num -lt 1 -or $num -gt $MaxNumber) {
                Write-InfoBox "Invalid number. Please enter 1-$MaxNumber" "Error"
                continue
            }

            return @($num)
        }

        Write-InfoBox "Invalid input. Please try again" "Error"
    }
}

function Invoke-Script {
    param(
        [string]$ScriptPath,
        [string]$ScriptName
    )

    Write-Host ""
    Write-InfoBox "Executing: $ScriptName" "Info"
    Write-Host ""

    $extension = [System.IO.Path]::GetExtension($ScriptPath)

    try {
        if ($extension -eq '.ps1') {
            & $ScriptPath
            $exitCode = $LASTEXITCODE
        }
        else {
            $process = Start-Process -FilePath $ScriptPath -Wait -PassThru -NoNewWindow
            $exitCode = $process.ExitCode
        }

        Write-Host ""
        if ($exitCode -eq 0 -or $null -eq $exitCode) {
            Write-InfoBox "Completed: $ScriptName" "Success"
        }
        else {
            Write-InfoBox "Completed with exit code $exitCode: $ScriptName" "Warning"
            $Script:ExitCode = $exitCode
        }
    }
    catch {
        Write-InfoBox "Failed to execute: $ScriptName" "Error"
        Write-InfoBox "Error: $($_.Exception.Message)" "Error"
        $Script:ExitCode = 1
    }
}

function Start-ScriptMenu {
    try {
        Clear-Host

        Write-Header "WinPE SCRIPT MENU"

        if (-not (Test-Path $ScriptsPath)) {
            Write-InfoBox "Scripts folder not found: $ScriptsPath" "Error"
            Write-Host ""
            Write-Host "Press any key to exit..." -ForegroundColor Gray
            $null = $Host.UI.RawUI.ReadKey('NoEcho,IncludeKeyDown')
            exit 1
        }

        Write-InfoBox "Scripts Folder: $ScriptsPath" "Info"

        $scripts = Get-ScriptFiles -Path $ScriptsPath

        if ($null -eq $scripts -or $scripts.Count -eq 0) {
            Write-Host ""
            Write-InfoBox "No scripts found" "Warning"
            Write-InfoBox "Supported: $($Script:SupportedExtensions -join ', ')" "Info"
            Write-Host ""
            Write-Host "Press any key to exit..." -ForegroundColor Gray
            $null = $Host.UI.RawUI.ReadKey('NoEcho,IncludeKeyDown')
            exit 1
        }

        Write-InfoBox "Found $($scripts.Count) script(s)" "Success"
        Write-Host ""
        Write-Separator
        Write-Host ""

        Show-ScriptMenu -Scripts $scripts

        $selections = Get-UserSelection -MaxNumber $scripts.Count

        if ($null -eq $selections) {
            Write-Host ""
            Write-InfoBox "Exiting..." "Info"
            Write-Host ""
            exit 0
        }

        Write-Host ""
        Write-Header "Executing Selected Scripts"

        foreach ($selection in $selections) {
            $script = $scripts[$selection - 1]
            Invoke-Script -ScriptPath $script.FullName -ScriptName $script.Name
        }

        Write-Host ""
        Write-Separator
        Write-InfoBox "All scripts completed" "Success"
        Write-Host ""
        Write-Host "Press any key to exit..." -ForegroundColor Gray
        $null = $Host.UI.RawUI.ReadKey('NoEcho,IncludeKeyDown')

        exit $Script:ExitCode
    }
    catch {
        Write-Host ""
        Write-InfoBox "FATAL ERROR: $($_.Exception.Message)" "Error"
        Write-Host ""
        Write-Host "Press any key to exit..." -ForegroundColor Gray
        $null = $Host.UI.RawUI.ReadKey('NoEcho,IncludeKeyDown')
        exit 1
    }
}

Start-ScriptMenu