# =========================================================
# MAX Windows Toolkit - 启动 PE 维护工具集
# 中文界面，统一风格
# =========================================================

$Common = "C:\Setup\windows\lib\common.ps1"
if (Test-Path $Common) { . $Common }

Write-Header "启动 PE 维护工具集"

Show-AdminWarn

# ---------------------------------------------------------
# 配置
# ---------------------------------------------------------

$PeToolBaseUrl = "$Global:RepoBase/tools-bin/pe"
$LocalPeToolDir = "C:\Setup\tools\pe"

New-Item -ItemType Directory -Path $LocalPeToolDir -Force | Out-Null

$PeTools = @(
    @{
        Id = 1
        Name = "BOOTICE"
        Type = "exe"
        FileName = "BOOTICE.exe"
        LocalPath = "BOOTICE.exe"
        Description = "引导/BCD/MBR/PBR 维护工具"
        RequireConfirm = $true
    },
    @{
        Id = 2
        Name = "CPU-Z"
        Type = "exe"
        FileName = "cpuz.exe"
        LocalPath = "cpuz.exe"
        Description = "CPU、主板和内存信息"
        RequireConfirm = $false
    },
    @{
        Id = 3
        Name = "Dism++"
        Type = "zip"
        FileName = "Dism++.zip"
        ExtractDir = "DismPlus"
        ExePattern = "Dism++*.exe"
        Description = "Windows 映像和系统维护工具"
        RequireConfirm = $true
    },
    @{
        Id = 4
        Name = "Everything"
        Type = "exe"
        FileName = "Everything.exe"
        LocalPath = "Everything.exe"
        Description = "快速文件搜索工具"
        RequireConfirm = $false
    },
    @{
        Id = 5
        Name = "WinNTSetup"
        Type = "exe"
        FileName = "WinNTSetup.exe"
        LocalPath = "WinNTSetup.exe"
        Description = "Windows 部署和安装工具"
        RequireConfirm = $true
    },
    @{
        Id = 6
        Name = "虚拟磁盘格式转换工具"
        Type = "exe"
        FileName = "VirtualDiskConverter.exe"
        LocalPath = "VirtualDiskConverter.exe"
        Description = "虚拟磁盘格式转换工具"
        RequireConfirm = $true
    },
    @{
        Id = 7
        Name = "Rufus"
        Type = "exe"
        FileName = "Rufus.exe"
        LocalPath = "Rufus.exe"
        Description = "启动盘制作工具"
        RequireConfirm = $true
    }
)

# ---------------------------------------------------------
# 辅助函数
# ---------------------------------------------------------

function Show-RiskWarning {
    param (
        [string]$ToolName
    )

    Write-Line
    Write-Warn "重要警告："
    Write-Warn "$ToolName 可能会修改引导记录、分区、映像、虚拟磁盘或系统部署设置。"
    Write-Err "错误操作可能导致系统无法启动或数据丢失。"
    Write-Info ""
    Write-Info "建议："
    Write-Info "1. 先备份重要数据"
    Write-Info "2. 仔细确认磁盘编号、分区编号和目标路径"
    Write-Info "3. 不要向需要恢复数据的磁盘写入"
    Write-Line
}

function Get-EncodedUrl {
    param (
        [string]$BaseUrl,
        [string]$FileName
    )

    $EncodedFileName = [System.Uri]::EscapeDataString($FileName)
    return "$BaseUrl/$EncodedFileName"
}

function Download-PeTool {
    param (
        [hashtable]$Tool
    )

    $FileName = $Tool.FileName
    $Url = Get-EncodedUrl -BaseUrl $PeToolBaseUrl -FileName $FileName
    $OutFile = Join-Path $LocalPeToolDir $FileName

    if (Test-Path $OutFile) {
        Write-OK "本地文件已存在: $FileName"
        return $true
    }

    Write-Info "正在下载: $FileName"

    if (Download-File -Url $Url -OutFile $OutFile) {
        try {
            Unblock-File -Path $OutFile -ErrorAction SilentlyContinue
        } catch {}
        return $true
    }

    Write-Err "下载失败: $FileName"
    return $false
}

function Update-PeTool {
    param (
        [hashtable]$Tool
    )

    $FileName = $Tool.FileName
    $OutFile = Join-Path $LocalPeToolDir $FileName

    if (Test-Path $OutFile) {
        Remove-Item $OutFile -Force -ErrorAction SilentlyContinue
    }

    if ($Tool.Type -eq "zip") {
        $ExtractPath = Join-Path $LocalPeToolDir $Tool.ExtractDir
        if (Test-Path $ExtractPath) {
            Remove-Item $ExtractPath -Recurse -Force -ErrorAction SilentlyContinue
        }
    }

    return Download-PeTool -Tool $Tool
}

function Get-ZipToolExe {
    param (
        [hashtable]$Tool
    )

    $ZipFile = Join-Path $LocalPeToolDir $Tool.FileName
    $ExtractPath = Join-Path $LocalPeToolDir $Tool.ExtractDir

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

    if (-not (Test-Path $ExtractPath)) {
        New-Item -ItemType Directory -Path $ExtractPath -Force | Out-Null

        Write-Info "正在解压: $($Tool.FileName)"

        try {
            Expand-Archive -Path $ZipFile -DestinationPath $ExtractPath -Force
            Get-ChildItem -Path $ExtractPath -Recurse -ErrorAction SilentlyContinue |
                Unblock-File -ErrorAction SilentlyContinue
        }
        catch {
            Write-Err "解压失败: $($Tool.FileName)"
            Write-Muted $_.Exception.Message
            return $null
        }
    }

    $Exe = Get-ChildItem -Path $ExtractPath -Filter $Tool.ExePattern -Recurse -ErrorAction SilentlyContinue |
        Select-Object -First 1

    if ($Exe) {
        return $Exe.FullName
    }

    return $null
}

function Get-PeToolExecutable {
    param (
        [hashtable]$Tool
    )

    if ($Tool.Type -eq "exe") {
        $ExePath = Join-Path $LocalPeToolDir $Tool.LocalPath
        if (Test-Path $ExePath) {
            return $ExePath
        }
        return $null
    }

    if ($Tool.Type -eq "zip") {
        return Get-ZipToolExe -Tool $Tool
    }

    return $null
}

function Start-PeTool {
    param (
        [hashtable]$Tool
    )

    if ($Tool.RequireConfirm) {
        Show-RiskWarning -ToolName $Tool.Name

        $Confirm = Read-Host "是否启动 $($Tool.Name)？输入 YES 继续"

        if ($Confirm -ne "YES") {
            Write-Warn "已取消"
            return
        }
    }

    if (-not (Download-PeTool -Tool $Tool)) {
        return
    }

    $ExePath = Get-PeToolExecutable -Tool $Tool

    if ([string]::IsNullOrWhiteSpace($ExePath)) {
        Write-Err "未找到 $($Tool.Name) 的可执行文件"
        Write-Warn "请检查工具包结构"
        return
    }

    Write-Info "正在启动: $($Tool.Name)"
    Write-Muted $ExePath

    try {
        Start-Process -FilePath $ExePath -WorkingDirectory (Split-Path $ExePath -Parent) -Verb RunAs
        Write-OK "$($Tool.Name) 已启动"
    }
    catch {
        Write-Err "启动失败: $($Tool.Name)"
        Write-Muted $_.Exception.Message
    }
}

function Download-AllPeTools {
    Write-Section "正在下载所有 PE 维护工具"

    foreach ($Tool in $PeTools) {
        Download-PeTool -Tool $Tool | Out-Null
    }

    Write-OK "全部下载完成"
}

function Update-AllPeTools {
    Write-Section "正在更新所有 PE 维护工具"

    foreach ($Tool in $PeTools) {
        Update-PeTool -Tool $Tool | Out-Null
    }

    Write-OK "全部更新完成"
}

function Show-PeToolList {
    Write-Info ""
    Write-Info "可用 PE 维护工具："
    Write-Info ""

    foreach ($Tool in $PeTools) {
        Write-Info "$($Tool.Id). $($Tool.Name) - $($Tool.Description)"
    }

    Write-Info ""
}

function Open-PeToolFolder {
    if (Test-Path $LocalPeToolDir) {
        Start-Process explorer.exe $LocalPeToolDir
    }
}

function Show-LocalDiskInfo {
    Write-Info ""
    Write-Info "物理磁盘："
    Get-Disk | Format-Table Number, FriendlyName, SerialNumber, BusType, PartitionStyle, OperationalStatus, Size -AutoSize

    Write-Info ""
    Write-Info "卷："
    Get-Volume | Format-Table DriveLetter, FileSystemLabel, FileSystem, DriveType, HealthStatus, SizeRemaining, Size -AutoSize

    Write-Info ""
    Write-Info "分区："
    Get-Partition | Format-Table DiskNumber, PartitionNumber, DriveLetter, Type, Size -AutoSize
}

# ---------------------------------------------------------
# 主菜单循环
# ---------------------------------------------------------

while ($true) {
    Clear-Host
    Write-Header "启动 PE 维护工具集"

    Write-Info "工具源站:"
    Write-Muted $PeToolBaseUrl
    Write-Info "本地路径:"
    Write-Muted $LocalPeToolDir
    Write-Info ""

    Show-PeToolList

    Write-Info "A. 下载全部工具"
    Write-Info "U. 更新全部工具"
    Write-Info "O. 打开本地 PE 工具文件夹"
    Write-Info "D. 显示本地磁盘信息"
    Write-Info "0. 退出"
    Write-Info ""

    $Choice = Read-Host "请选择工具编号或选项"

    if ($Choice -match '^[0]$') {
        break
    }

    if ($Choice -match '^[Aa]$') {
        Download-AllPeTools
        Pause-Max
        continue
    }

    if ($Choice -match '^[Uu]$') {
        Update-AllPeTools
        Pause-Max
        continue
    }

    if ($Choice -match '^[Oo]$') {
        Open-PeToolFolder
        Pause-Max
        continue
    }

    if ($Choice -match '^[Dd]$') {
        Show-LocalDiskInfo
        Pause-Max
        continue
    }

    if ($Choice -match '^\d+$') {
        $Id = [int]$Choice
        $Tool = $PeTools | Where-Object { $_.Id -eq $Id } | Select-Object -First 1

        if ($Tool) {
            Start-PeTool -Tool $Tool
            Pause-Max
        } else {
            Write-Err "无效的工具编号: $Choice"
            Pause-Max
        }

        continue
    }

    Write-Err "无效选项"
    Pause-Max
}