# Download Qwen Image 2.1 GGUF + companions into the live Comfy Desktop Shared models tree.
# Also ensures city96/ComfyUI-GGUF is present on the ComfyUI (1) install.
#
# Example:
#   powershell -ExecutionPolicy Bypass -File scripts\setup-qwen21.ps1
#   powershell -File scripts\setup-qwen21.ps1 -SkipDownload
#   powershell -File scripts\setup-qwen21.ps1 -ComfyRoot "C:\...\ComfyUI (1)\ComfyUI"

[CmdletBinding()]
param(
  [string]$ComfyRoot = '',
  [string]$ModelsRoot = '',
  [switch]$SkipDownload,
  [switch]$SkipNode,
  [switch]$Force
)

$ErrorActionPreference = 'Stop'
$ProgressPreference = 'Continue'

function Write-Step([string]$Message) {
  Write-Host ""
  Write-Host "==> $Message" -ForegroundColor Cyan
}

function Write-Ok([string]$Message) {
  Write-Host "    $Message" -ForegroundColor Green
}

function Write-Warn([string]$Message) {
  Write-Host "    warn  $Message" -ForegroundColor Yellow
}

function Resolve-ComfyRoot([string]$Hint) {
  $candidates = @(
    $Hint,
    $env:COMFY_ROOT,
    $env:COMFYUI_ROOT,
    (Join-Path $env:LOCALAPPDATA 'Comfy-Desktop\ComfyUI-Installs\ComfyUI (1)\ComfyUI'),
    (Join-Path $env:LOCALAPPDATA 'Comfy-Desktop\ComfyUI-Installs\ComfyUI\ComfyUI'),
    'C:\ComfyUI',
    (Join-Path $env:USERPROFILE 'ComfyUI')
  ) | Where-Object { $_ } | ForEach-Object { $_.TrimEnd('\', '/') }

  foreach ($path in $candidates) {
    if (Test-Path (Join-Path $path 'main.py')) { return (Resolve-Path $path).Path }
    if (Test-Path (Join-Path $path 'ComfyUI\main.py')) { return (Resolve-Path (Join-Path $path 'ComfyUI')).Path }
  }
  throw "Could not find ComfyUI root (main.py). Pass -ComfyRoot."
}

function Resolve-ModelsRoot([string]$Hint, [string]$ComfyRootPath) {
  $shared = Join-Path $env:LOCALAPPDATA 'Comfy-Desktop\ComfyUI-Shared\models'
  $candidates = @(
    $Hint,
    $env:COMFY_MODELS_ROOT,
    $shared,
    (Join-Path $ComfyRootPath 'models')
  ) | Where-Object { $_ } | ForEach-Object { $_.TrimEnd('\', '/') }

  foreach ($path in $candidates) {
    if (Test-Path $path) { return (Resolve-Path $path).Path }
  }
  New-Item -ItemType Directory -Force -Path $shared | Out-Null
  return (Resolve-Path $shared).Path
}

function Find-HfCli {
  foreach ($name in @('hf', 'huggingface-cli')) {
    $cmd = Get-Command $name -ErrorAction SilentlyContinue
    if ($cmd) { return $cmd.Source }
  }
  throw "Hugging Face CLI not found (hf / huggingface-cli). Install with: pip install -U huggingface_hub"
}

function Ensure-Dir([string]$Path) {
  if (-not (Test-Path $Path)) {
    New-Item -ItemType Directory -Force -Path $Path | Out-Null
  }
}

function Assert-ModelFile([string]$Path, [long]$MinBytes) {
  if (-not (Test-Path $Path)) { throw "Missing model file: $Path" }
  $item = Get-Item $Path
  if ($item.Length -lt $MinBytes) {
    throw "Suspicious size for $Path ($($item.Length) bytes). Expected >= $MinBytes. Delete and re-run."
  }
  Write-Ok ("{0}  {1:N0} bytes" -f $item.FullName, $item.Length)
}

function Flatten-IfNested([string]$DestDir, [string]$FileName) {
  $nested = Get-ChildItem $DestDir -Recurse -File -Filter $FileName -ErrorAction SilentlyContinue |
    Where-Object { $_.DirectoryName -ne $DestDir } |
    Select-Object -First 1
  if ($nested) {
    $target = Join-Path $DestDir $FileName
    if ((Test-Path $target) -and -not $Force) {
      Write-Warn "Nested copy found but destination already exists: $target"
      return
    }
    Move-Item -Force $nested.FullName $target
    Write-Ok "Flattened $($nested.FullName) -> $target"
  }
}

function Invoke-HfFile {
  param(
    [string]$Cli,
    [string]$Repo,
    [string]$File,
    [string]$LocalDir
  )
  Ensure-Dir $LocalDir
  $dest = Join-Path $LocalDir (Split-Path $File -Leaf)
  if ((Test-Path $dest) -and -not $Force) {
    $len = (Get-Item $dest).Length
    if ($len -gt 1MB) {
      Write-Ok "Already present: $dest ($len bytes)"
      return
    }
  }
  Write-Step "hf download $Repo $File"
  & $Cli download $Repo $File --local-dir $LocalDir
  if ($LASTEXITCODE -ne 0) { throw "hf download failed for $Repo / $File (exit $LASTEXITCODE)" }
  Flatten-IfNested -DestDir $LocalDir -FileName (Split-Path $File -Leaf)
}

# --- resolve paths ---
Write-Step "Resolve ComfyUI root"
$root = Resolve-ComfyRoot $ComfyRoot
Write-Ok $root

Write-Step "Resolve models root (Desktop Shared preferred)"
$models = Resolve-ModelsRoot $ModelsRoot $root
Write-Ok $models

$Diff = Join-Path $models 'diffusion_models'
$TE   = Join-Path $models 'text_encoders'
$VAE  = Join-Path $models 'vae'
Ensure-Dir $Diff
Ensure-Dir $TE
Ensure-Dir $VAE

if (-not $SkipDownload) {
  $cli = Find-HfCli
  Write-Ok "HF CLI: $cli"

  # DiT GGUF only (no Q6/Q5/Q4, no full bf16 DiT)
  Invoke-HfFile -Cli $cli -Repo 'abenzerps/Qwen-Image-2.1-Uncensored-GGUF' -File 'qwen-image-2.1-Q8_0.gguf' -LocalDir $Diff

  # TE + VAE companions from Comfy-Org (INT8 TE for 16 GB VRAM)
  Invoke-HfFile -Cli $cli -Repo 'Comfy-Org/Qwen-Image-2.1' -File 'text_encoders/qwen3vl_8b_int8_convrot.safetensors' -LocalDir $TE
  Invoke-HfFile -Cli $cli -Repo 'Comfy-Org/Qwen-Image-2.1' -File 'vae/qwen_image_2.1_vae_bf16.safetensors' -LocalDir $VAE
}

Write-Step "Verify weights"
Assert-ModelFile (Join-Path $Diff 'qwen-image-2.1-Q8_0.gguf') 7000000000
Assert-ModelFile (Join-Path $TE 'qwen3vl_8b_int8_convrot.safetensors') 1000000000
Assert-ModelFile (Join-Path $VAE 'qwen_image_2.1_vae_bf16.safetensors') 100000000

if (-not $SkipNode) {
  Write-Step "city96 ComfyUI-GGUF"
  $nodes = Join-Path $root 'custom_nodes\ComfyUI-GGUF'
  if (-not (Test-Path $nodes)) {
    git clone https://github.com/city96/ComfyUI-GGUF $nodes
    if ($LASTEXITCODE -ne 0) { throw "git clone ComfyUI-GGUF failed" }
    Write-Ok "Cloned $nodes"
  } else {
    Write-Ok "Already present: $nodes"
  }

  $py = Join-Path $root '.venv\Scripts\python.exe'
  if (-not (Test-Path $py)) {
    $py = Join-Path (Split-Path $root) 'standalone-env\python.exe'
  }
  if (Test-Path $py) {
    Write-Step "Ensure gguf package in Comfy python ($py)"
    & $py -m pip install -q gguf
    if ($LASTEXITCODE -ne 0) { Write-Warn "pip install gguf returned $LASTEXITCODE" }
    else { Write-Ok "gguf importable" }
  } else {
    Write-Warn "Comfy python not found; install gguf into the interpreter that runs main.py"
  }
}

Write-Host ""
Write-Host "Done. Restart Comfy only when /status shows gpu.busy=false." -ForegroundColor Cyan
Write-Host "ComfyRoot=$root"
Write-Host "ModelsRoot=$models"