#requires -Version 5.1
[CmdletBinding()]
param(
    [ValidateSet('auto', 'full', 'lite', 'finish', 'rtx_creator', 'rtx_low_vram', 'cuda_universal', 'non_nvidia_experimental', 'media_finish_only')]
    [string]$Profile = 'auto',
    [switch]$SkipModels,
    [switch]$DryRun,
    [switch]$NoLaunch
)

Set-StrictMode -Version Latest
$ErrorActionPreference = 'Stop'
[Net.ServicePointManager]::SecurityProtocol = [Net.SecurityProtocolType]::Tls12

$InstallerRoot = [IO.Path]::GetFullPath($PSScriptRoot)
$ProjectRoot = [IO.Path]::GetFullPath((Join-Path $InstallerRoot '..'))
$RuntimeRoot = Join-Path $ProjectRoot '.runtime\python'
$Python = Join-Path $RuntimeRoot 'python.exe'
$CacheRoot = Join-Path $ProjectRoot '.install-cache'
$LogFile = Join-Path $ProjectRoot 'install_h3_cn.log'
$Requirements = Join-Path $InstallerRoot 'requirements-cn.txt'
$StudioRequirements = Join-Path $InstallerRoot 'requirements-studio.txt'
$ComfyRoot = Join-Path $ProjectRoot 'ComfyUI'

function Write-Stage([int]$Number, [string]$Message) {
    $line = "[$Number/5] $Message"
    Write-Host "`n$line" -ForegroundColor Cyan
    Add-Content -LiteralPath $LogFile -Value "$(Get-Date -Format s) $line" -Encoding UTF8
}

function Write-Note([string]$Message) {
    Write-Host "  $Message" -ForegroundColor DarkGray
    Add-Content -LiteralPath $LogFile -Value "$(Get-Date -Format s) $Message" -Encoding UTF8
}

function Invoke-Checked([string]$Label, [scriptblock]$Action) {
    Write-Note $Label
    & $Action
    if ($LASTEXITCODE -ne 0) { throw "$Label 失败，退出码 $LASTEXITCODE" }
}

function Get-RemoteFile([string[]]$Urls, [string]$Destination) {
    if (Test-Path -LiteralPath $Destination -PathType Leaf) {
        Write-Note "已存在，跳过下载：$([IO.Path]::GetFileName($Destination))"
        return
    }
    $parent = Split-Path -Parent $Destination
    New-Item -ItemType Directory -Path $parent -Force | Out-Null
    $lastError = $null
    foreach ($url in $Urls) {
        try {
            Write-Note "国内源下载：$url"
            if (Get-Command Start-BitsTransfer -ErrorAction SilentlyContinue) {
                Start-BitsTransfer -Source $url -Destination $Destination -DisplayName 'H3 Studio 安装文件' -ErrorAction Stop
            } else {
                Invoke-WebRequest -Uri $url -OutFile $Destination -UseBasicParsing -ErrorAction Stop
            }
            if ((Get-Item -LiteralPath $Destination).Length -gt 1MB) { return }
            throw '下载文件尺寸异常'
        } catch {
            $lastError = $_
            if (Test-Path -LiteralPath $Destination -PathType Leaf) { Remove-Item -LiteralPath $Destination -Force -ErrorAction SilentlyContinue }
        }
    }
    throw "所有国内下载源均失败：$lastError"
}

function Get-GpuPlan {
    $gpuName = '未检测到 NVIDIA GPU'
    $gpuVendor = 'cpu'
    $vramGb = 0.0
    $smi = Get-Command nvidia-smi.exe -ErrorAction SilentlyContinue
    if ($smi) {
        $raw = & $smi.Source '--query-gpu=name,memory.total' '--format=csv,noheader,nounits' 2>$null | Select-Object -First 1
        if ($raw -match '^(.+),\s*([0-9.]+)$') {
            $gpuName = $Matches[1].Trim()
            $gpuVendor = 'nvidia'
            $vramGb = [math]::Round(([double]$Matches[2] / 1024), 1)
        }
    }
    if ($gpuVendor -eq 'cpu') {
        try {
            $controller = Get-CimInstance Win32_VideoController -ErrorAction Stop |
                Where-Object { $_.Name -notmatch 'Microsoft Basic' } |
                Sort-Object AdapterRAM -Descending |
                Select-Object -First 1
            if ($controller) {
                $gpuName = [string]$controller.Name
                if ($gpuName -match 'AMD|Radeon') { $gpuVendor = 'amd' }
                elseif ($gpuName -match 'Intel|Arc') { $gpuVendor = 'intel' }
                elseif ($gpuName -match 'NVIDIA') { $gpuVendor = 'nvidia' }
                if ($controller.AdapterRAM) { $vramGb = [math]::Round(([double]$controller.AdapterRAM / 1GB), 1) }
            }
        } catch {
            Write-Note 'Windows 未返回显卡信息，将使用通用媒体终修规格。'
        }
    }
    $selected = $Profile
    $aliases = @{ full = 'rtx_creator'; lite = 'rtx_low_vram'; finish = 'media_finish_only' }
    if ($aliases.ContainsKey($selected)) { $selected = $aliases[$selected] }
    if ($selected -eq 'auto') {
        if ($gpuName -match 'RTX' -and $vramGb -ge 23.5) { $selected = 'rtx_creator' }
        elseif ($gpuName -match 'RTX' -and $vramGb -ge 15.5) { $selected = 'rtx_low_vram' }
        elseif ($gpuVendor -eq 'nvidia' -and $vramGb -ge 23.5) { $selected = 'cuda_universal' }
        elseif ($gpuVendor -in @('amd', 'intel') -and $vramGb -ge 15.5) { $selected = 'non_nvidia_experimental' }
        else { $selected = 'media_finish_only' }
    }
    [pscustomobject]@{ Name = $gpuName; Vendor = $gpuVendor; VramGb = $vramGb; Profile = $selected }
}

function Install-PortablePython {
    if (Test-Path -LiteralPath $Python -PathType Leaf) {
        Write-Note '独立 Python 运行时已存在，跳过。'
        return
    }
    $offline = @(
        (Join-Path $ProjectRoot 'packages\H3-Studio-Python312-Runtime.zip'),
        (Join-Path $ProjectRoot 'packages\H3-Studio-Python312-CUDA130-Runtime.zip')
    ) | Where-Object { Test-Path -LiteralPath $_ -PathType Leaf } | Select-Object -First 1
    if ($offline) {
        Write-Note '发现离线运行时包，正在解压。'
        New-Item -ItemType Directory -Path (Split-Path -Parent $RuntimeRoot) -Force | Out-Null
        Expand-Archive -LiteralPath $offline -DestinationPath (Split-Path -Parent $RuntimeRoot) -Force
        if (-not (Test-Path -LiteralPath $Python -PathType Leaf)) { throw '离线运行时包结构不正确' }
        return
    }
    $installer = Join-Path $CacheRoot 'python-3.12.10-amd64.exe'
    Get-RemoteFile -Urls @(
        'https://mirrors.huaweicloud.com/python/3.12.10/python-3.12.10-amd64.exe',
        'https://mirrors.aliyun.com/python-release/windows/python-3.12.10-amd64.exe'
    ) -Destination $installer
    New-Item -ItemType Directory -Path $RuntimeRoot -Force | Out-Null
    $arguments = "/quiet InstallAllUsers=0 TargetDir=`"$RuntimeRoot`" Include_pip=1 Include_launcher=0 Include_test=0 Include_doc=0 Include_tcltk=0 PrependPath=0 Shortcuts=0"
    $process = Start-Process -FilePath $installer -ArgumentList $arguments -Wait -PassThru
    if ($process.ExitCode -ne 0 -or -not (Test-Path -LiteralPath $Python -PathType Leaf)) {
        throw "独立 Python 安装失败，退出码 $($process.ExitCode)"
    }
}

function Install-PythonPackages([string]$SelectedProfile) {
    $tuna = 'https://pypi.tuna.tsinghua.edu.cn/simple'
    Invoke-Checked '更新 pip 与基础打包工具' { & $Python -m pip install --upgrade pip setuptools wheel -i $tuna --disable-pip-version-check }
    $cudaProfiles = @('rtx_creator', 'rtx_low_vram', 'cuda_universal')
    if ($SelectedProfile -in $cudaProfiles) {
        Invoke-Checked '从阿里云安装 PyTorch CUDA 13.0' {
            & $Python -m pip install 'torch==2.10.0+cu130' 'torchvision==0.25.0+cu130' 'torchaudio==2.10.0+cu130' --index-url 'https://mirrors.aliyun.com/pytorch-wheels/cu130' --disable-pip-version-check
        }
    } elseif ($SelectedProfile -eq 'non_nvidia_experimental') {
        Write-Warning 'AMD / Intel 的 H3 计算后端取决于本机驱动。先安装通用组件，之后由部署向导检测 ROCm、XPU 或 DirectML。'
    }
    $selectedRequirements = if ($SelectedProfile -eq 'media_finish_only') { $StudioRequirements } else { $Requirements }
    $dependencyLabel = if ($SelectedProfile -eq 'media_finish_only') { '从清华 TUNA 安装轻量 Studio 终修依赖' } else { '从清华 TUNA 安装 H3 Studio 与 ComfyUI 依赖' }
    Invoke-Checked $dependencyLabel { & $Python -m pip install -r $selectedRequirements -i $tuna --disable-pip-version-check }
    if ($SelectedProfile -in $cudaProfiles) {
        & $Python -m pip install 'triton-windows>=3.5' -i $tuna --disable-pip-version-check
        if ($LASTEXITCODE -ne 0) { Write-Warning 'Triton 安装失败；首次运行将使用 PyTorch Attention，不影响基础生成。' }
    }
}

function Copy-ModelFile([string]$DownloadRoot, [string]$RelativePath, [string]$TargetFolder) {
    $name = [IO.Path]::GetFileName($RelativePath)
    $target = Join-Path $TargetFolder $name
    if (Test-Path -LiteralPath $target -PathType Leaf) {
        Write-Note "模型已存在：$name"
        return
    }
    $source = Join-Path $DownloadRoot ($RelativePath -replace '/', '\')
    if (-not (Test-Path -LiteralPath $source -PathType Leaf)) {
        $source = Get-ChildItem -LiteralPath $DownloadRoot -Recurse -File -Filter $name | Select-Object -First 1 -ExpandProperty FullName
    }
    if (-not $source -or -not (Test-Path -LiteralPath $source -PathType Leaf)) { throw "模型下载后未找到：$name" }
    New-Item -ItemType Directory -Path $TargetFolder -Force | Out-Null
    Copy-Item -LiteralPath $source -Destination $target -Force
}

function Install-H3Models([string]$SelectedProfile) {
    if ($SkipModels -or $SelectedProfile -eq 'media_finish_only') {
        Write-Note '当前方案跳过 H3 大模型下载。'
        return
    }
    $downloadRoot = Join-Path $CacheRoot 'modelscope-h3'
    $files = @(
        'diffusion_models/minimax_h3_fl2va_pruned_int8_convrot.safetensors',
        'text_encoders/qwen3vl_32b_minimax_h3_nvfp4_awq.safetensors',
        'vae/minimax_h3_video_vae_fp16.safetensors',
        'vae/minimax_h3_audio_vae_fp32.safetensors'
    )
    $files += 'diffusion_models/minimax_h3_ref2va_pruned_int8_convrot.safetensors'
    $modelscope = Join-Path $RuntimeRoot 'Scripts\modelscope.exe'
    if (-not (Test-Path -LiteralPath $modelscope -PathType Leaf)) { throw 'ModelScope CLI 未正确安装' }
    New-Item -ItemType Directory -Path $downloadRoot -Force | Out-Null
    Write-Note '从魔搭 ModelScope 下载 H3 权重；已下载文件会自动续传。'
    & $modelscope download --model 'Comfy-Org/MiniMax-H3' --local_dir $downloadRoot @files
    if ($LASTEXITCODE -ne 0) { throw 'MiniMax H3 模型下载失败，可重新运行安装器续传。' }
    foreach ($relative in $files) {
        $folder = if ($relative.StartsWith('diffusion_models/')) { Join-Path $ComfyRoot 'models\diffusion_models' } elseif ($relative.StartsWith('text_encoders/')) { Join-Path $ComfyRoot 'models\text_encoders' } else { Join-Path $ComfyRoot 'models\vae' }
        Copy-ModelFile -DownloadRoot $downloadRoot -RelativePath $relative -TargetFolder $folder
    }
    $turboTarget = Join-Path $ComfyRoot 'models\loras\minimax_h3_turbo_v4_step600_ema.safetensors'
    if (-not (Test-Path -LiteralPath $turboTarget -PathType Leaf)) {
        New-Item -ItemType Directory -Path (Split-Path -Parent $turboTarget) -Force | Out-Null
        $env:HF_ENDPOINT = 'https://hf-mirror.com'
        $code = "from huggingface_hub import hf_hub_download; hf_hub_download(repo_id='larryvrh/MiniMax-H3-Turbo-Lora', filename='minimax_h3_turbo_v4_step600_ema.safetensors', local_dir=r'$((Split-Path -Parent $turboTarget).Replace("'", "''"))')"
        & $Python -c $code
        if ($LASTEXITCODE -ne 0) { throw 'Turbo LoRA 国内镜像下载失败；请稍后重新运行安装器。' }
    }
}

function Save-Profile([string]$SelectedProfile) {
    $value = [ordered]@{
        version = 2
        profile_id = $SelectedProfile
        hardware_override = @{}
        configured_at = (Get-Date).ToUniversalTime().ToString('o')
        installed_by = 'h3_cn_oneclick'
    }
    $config = Join-Path $ProjectRoot 'studio\deployment_config.json'
    $value | ConvertTo-Json -Depth 5 | Set-Content -LiteralPath $config -Encoding UTF8
}

try {
    Set-Content -LiteralPath $LogFile -Value "H3 Studio CN installer started $(Get-Date -Format s)" -Encoding UTF8
    Write-Stage 1 '检测硬件并选择最省事的安装方案'
    $plan = Get-GpuPlan
    $driveName = [IO.Path]::GetPathRoot($ProjectRoot).Substring(0, 1)
    $freeGb = [math]::Round(([IO.DriveInfo]::new($driveName)).AvailableFreeSpace / 1GB, 1)
    Write-Host "  GPU：$($plan.Name) · $($plan.VramGb) GB"
    Write-Host "  自动方案：$($plan.Profile) · 可用磁盘 $freeGb GB" -ForegroundColor Green
    if ($plan.Profile -ne 'media_finish_only' -and $freeGb -lt 100) { throw '磁盘空间不足：H3 完整安装建议至少 100GB 可用空间。' }
    if (-not (Test-Path -LiteralPath (Join-Path $ComfyRoot 'main.py') -PathType Leaf)) { throw '安装包不完整：缺少 ComfyUI\main.py' }
    if ($DryRun) {
        [pscustomobject]@{ gpu = $plan.Name; vram_gb = $plan.VramGb; profile = $plan.Profile; disk_free_gb = $freeGb; model_source = 'ModelScope China'; python_source = 'Huawei/Aliyun'; package_source = 'TUNA/Aliyun' } | ConvertTo-Json
        exit 0
    }
    if ($plan.Profile -ne 'media_finish_only' -and -not $SkipModels) {
        Write-Host "`n将下载 MiniMax H3 社区许可模型。启用完整生成、R2V 与重绘约需 60GB。" -ForegroundColor Yellow
        Write-Host '模型许可：https://www.modelscope.cn/models/Comfy-Org/MiniMax-H3/'
        if ((Read-Host '输入 Y 表示你已阅读并同意模型许可') -notmatch '^[Yy]$') { throw '用户取消模型下载。' }
    }
    Write-Stage 2 '准备独立 Python 运行时（不修改系统环境）'
    Install-PortablePython
    Write-Stage 3 '从国内镜像安装运行依赖'
    Install-PythonPackages -SelectedProfile $plan.Profile
    Write-Stage 4 '下载并放置 MiniMax H3 模型'
    Install-H3Models -SelectedProfile $plan.Profile
    Write-Stage 5 '保存配置并启动 H3 Studio'
    Save-Profile -SelectedProfile $plan.Profile
    if ($plan.Profile -eq 'media_finish_only') {
        Invoke-Checked '运行轻量终修环境自检' { & $Python -c "import aiohttp,av,numpy,PIL; print('H3 Studio finish runtime OK')" }
    } else {
        Invoke-Checked '运行 H3 环境自检' { & $Python -c "import aiohttp,av,numpy,cv2,modelscope,torch; print('H3 runtime OK')" }
    }
    Write-Host "`n安装完成。以后只需双击 start_h3_studio.cmd。" -ForegroundColor Green
    if (-not $NoLaunch) { Start-Process -FilePath (Join-Path $ProjectRoot 'start_h3_studio.cmd') -WorkingDirectory $ProjectRoot }
    exit 0
} catch {
    $message = $_.Exception.Message
    Add-Content -LiteralPath $LogFile -Value "$(Get-Date -Format s) ERROR $message" -Encoding UTF8
    Write-Host "`n安装未完成：$message" -ForegroundColor Red
    Write-Host '修复网络或磁盘问题后再次双击安装器，会跳过已完成部分。' -ForegroundColor Yellow
    exit 1
}
