# =========================================================
# MAX Windows Toolkit - 使用 DiskGenius 管理磁盘
# 中文界面，统一风格
# =========================================================

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

Write-Header "使用 DiskGenius 管理磁盘"

Show-AdminWarn

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

$RepoBase = "https://dl.unvmax.com/windows"
$ToolsBaseUrl = "$RepoBase/tools-bin"

$LocalRoot = "C:\Setup\windows"
$ToolsDir  = Join-Path $LocalRoot "tools-bin"
$DiskGeniusDir = Join-Path $ToolsDir "DiskGenius"

$DiskGeniusExeName = "DiskGenius.exe"

$DiskGeniusExeUrl  = "$ToolsBaseUrl/$DiskGeniusExeName"
$DiskGeniusExePath = Join-Path $DiskGeniusDir $DiskGeniusExeName

New-Item -ItemType Directory -Path $ToolsDir -Force | Out-Null
New-Item -ItemType Directory -Path $DiskGeniusDir -Force | Out-Null

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

function Show-DiskGeniusWarning {
    Write-Line
    Write-Warn "重要警告："
    Write-Warn "DiskGenius 可以修改分区、格式化磁盘、克隆磁盘、恢复数据和重建分区表。"
    Write-Err "错误操作可能导致数据丢失！"
    Write-Info ""
    Write-Info "安全提示："
    Write-Info "  - 使用 DiskGenius 前请备份重要数据"
    Write-Info "  - 仔细确认磁盘编号和分区编号"
    Write-Info "  - 不要向需要恢复数据的磁盘写入"
    Write-Info "  - 恢复的文件请保存到另一块物理磁盘"
    Write-Info "  - 在确认目标磁盘前，避免格式化、重建分区表或克隆"
    Write-Line
}

function Get-DiskGeniusExePath {
    if (Test-Path $DiskGeniusExePath) {
        return $DiskGeniusExePath
    }

    $Found = Get-ChildItem -Path $DiskGeniusDir `
        -Filter "DiskGenius.exe" `
        -Recurse `
        -File `
        -ErrorAction SilentlyContinue |
        Select-Object -First 1

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

    $Found2 = Get-ChildItem -Path $ToolsDir `
        -Filter "DiskGenius.exe" `
        -Recurse `
        -File `
        -ErrorAction SilentlyContinue |
        Select-Object -First 1

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

    return $null
}

function Download-OrUpdateDiskGenius {
    Write-Section "下载 / 更新 DiskGenius"
    Write-Muted "URL : $DiskGeniusExeUrl"
    Write-Muted "路径: $DiskGeniusExePath"

    if (-not (Download-File -Url $DiskGeniusExeUrl -OutFile $DiskGeniusExePath)) {
        Write-Err "下载 DiskGenius 失败"
        return $false
    }

    try {
        Unblock-File -Path $DiskGeniusExePath -ErrorAction SilentlyContinue
    } catch {}

    if (Test-Path $DiskGeniusExePath) {
        Write-OK "DiskGenius 已就绪"
        Write-Muted "可执行文件: $DiskGeniusExePath"
        return $true
    }

    Write-Err "下载后未找到 DiskGenius.exe"
    return $false
}

function Ensure-DiskGenius {
    $Exe = Get-DiskGeniusExePath

    if ($Exe) {
        return $true
    }

    Write-Warn "本地未找到 DiskGenius"
    Write-Muted "期望路径: $DiskGeniusExePath"
    Write-Info ""

    $Confirm = Read-Host "现在下载？输入 YES 继续"

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

    return Download-OrUpdateDiskGenius
}

function Start-DiskGenius {
    if (-not (Ensure-DiskGenius)) {
        return
    }

    $Exe = Get-DiskGeniusExePath

    if (-not $Exe) {
        Write-Err "未找到 DiskGenius.exe"
        return
    }

    Write-Info "正在启动 DiskGenius..."
    Write-Muted "路径: $Exe"

    try {
        if (Test-IsAdmin) {
            Start-Process -FilePath $Exe -WorkingDirectory (Split-Path $Exe -Parent)
        } else {
            Write-Warn "当前会话不是管理员，尝试以管理员身份启动..."
            Start-Process `
                -FilePath $Exe `
                -WorkingDirectory (Split-Path $Exe -Parent) `
                -Verb RunAs
        }

        Write-OK "DiskGenius 已启动"
    }
    catch {
        Write-Err "启动 DiskGenius 失败"
        Write-Muted $_.Exception.Message
    }
}

function Open-ToolsFolder {
    New-Item -ItemType Directory -Path $ToolsDir -Force | Out-Null
    Start-Process explorer.exe $ToolsDir
}

function Open-DiskGeniusFolder {
    New-Item -ItemType Directory -Path $DiskGeniusDir -Force | Out-Null
    Start-Process explorer.exe $DiskGeniusDir
}

function Create-DiskGeniusShortcuts {
    if (-not (Ensure-DiskGenius)) {
        return
    }

    $Exe = Get-DiskGeniusExePath

    if (-not $Exe) {
        Write-Err "未找到 DiskGenius.exe"
        return
    }

    try {
        $Desktop = [Environment]::GetFolderPath("Desktop")
        $StartMenu = Join-Path $env:ProgramData "Microsoft\Windows\Start Menu\Programs"

        $ShortcutTargets = @(
            @{
                Path = Join-Path $Desktop "DiskGenius.lnk"
            },
            @{
                Path = Join-Path $StartMenu "DiskGenius.lnk"
            }
        )

        $Shell = New-Object -ComObject WScript.Shell

        foreach ($Target in $ShortcutTargets) {
            $Shortcut = $Shell.CreateShortcut($Target.Path)
            $Shortcut.TargetPath = $Exe
            $Shortcut.WorkingDirectory = Split-Path $Exe -Parent
            $Shortcut.IconLocation = $Exe
            $Shortcut.Save()

            Write-OK "快捷方式已创建: $($Target.Path)"
        }
    }
    catch {
        Write-Err "创建快捷方式失败"
        Write-Muted $_.Exception.Message
    }
}

function Show-LocalDiskInformation {
    Write-Info ""
    Write-Info "本地磁盘信息："
    Write-Info ""

    try {
        Write-Info "物理磁盘："

        Get-Disk |
            Select-Object Number, FriendlyName, SerialNumber, BusType, PartitionStyle, OperationalStatus, Size |
            Format-Table -AutoSize

        Write-Info ""
        Write-Info "卷："

        Get-Volume |
            Select-Object DriveLetter, FileSystemLabel, FileSystem, DriveType, HealthStatus, SizeRemaining, Size |
            Format-Table -AutoSize
    }
    catch {
        Write-Warn "PowerShell cmdlet 获取磁盘信息失败，尝试 WMIC 备用方案..."
        Write-Info ""

        try {
            wmic diskdrive get Index,Model,SerialNumber,Size,Status
            Write-Info ""
            wmic logicaldisk get DeviceID,VolumeName,FileSystem,FreeSpace,Size
        }
        catch {
            Write-Err "获取磁盘信息失败"
            Write-Muted $_.Exception.Message
        }
    }
}

function Open-WindowsDiskManagement {
    Write-Info "正在打开 Windows 磁盘管理..."
    try {
        Start-Process "diskmgmt.msc"
    }
    catch {
        Write-Err "打开磁盘管理失败"
        Write-Muted $_.Exception.Message
    }
}

function Open-WindowsStorageSettings {
    Write-Info "正在打开 Windows 存储设置..."
    try {
        Start-Process "ms-settings:storagesense"
    }
    catch {
        Write-Err "打开存储设置失败"
        Write-Muted $_.Exception.Message
    }
}

function Remove-LocalDiskGenius {
    Write-Warn "此操作将删除本地 DiskGenius 文件："
    Write-Muted $DiskGeniusDir
    Write-Info ""

    $Confirm = Read-Host "确认删除？输入 DELETE 确认"

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

    try {
        Get-Process -Name "DiskGenius" -ErrorAction SilentlyContinue |
            Stop-Process -Force -ErrorAction SilentlyContinue

        if (Test-Path $DiskGeniusDir) {
            Remove-Item $DiskGeniusDir -Recurse -Force -ErrorAction SilentlyContinue
            Write-OK "DiskGenius 文件夹已删除"
        } else {
            Write-Warn "DiskGenius 文件夹不存在"
        }
    }
    catch {
        Write-Err "删除 DiskGenius 失败"
        Write-Muted $_.Exception.Message
    }
}

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

:diskLoop while ($true) {
    Clear-Host
    Write-Header "使用 DiskGenius 管理磁盘"

    Show-DiskGeniusWarning

    Write-Info "1. 启动 DiskGenius"
    Write-Info "2. 下载 / 更新 DiskGenius"
    Write-Info "3. 打开工具文件夹"
    Write-Info "4. 打开 DiskGenius 文件夹"
    Write-Info "5. 创建桌面和开始菜单快捷方式"
    Write-Info "6. 显示本地磁盘信息"
    Write-Info "7. 打开 Windows 磁盘管理"
    Write-Info "8. 打开 Windows 存储设置"
    Write-Info "9. 删除本地 DiskGenius"
    Write-Info "0. 退出"
    Write-Info ""

    $Choice = Read-Host "请选择"

    switch ($Choice) {
        "1" {
            Start-DiskGenius
            Pause-Max
        }
        "2" {
            Download-OrUpdateDiskGenius | Out-Null
            Pause-Max
        }
        "3" {
            Open-ToolsFolder
            Pause-Max
        }
        "4" {
            Open-DiskGeniusFolder
            Pause-Max
        }
        "5" {
            Create-DiskGeniusShortcuts
            Pause-Max
        }
        "6" {
            Show-LocalDiskInformation
            Pause-Max
        }
        "7" {
            Open-WindowsDiskManagement
            Pause-Max
        }
        "8" {
            Open-WindowsStorageSettings
            Pause-Max
        }
        "9" {
            Remove-LocalDiskGenius
            Pause-Max
        }
        "0" { break :diskLoop }
        default {
            Write-Err "无效选项"
            Pause-Max
        }
    }
}