# Download Qwen2.5-VL NSFW Caption V4 GGUF (Q5_K_M + mmproj only) for Studio 2 caption jobs.
# Does NOT pull the full safetensors repo (~16.6GB).
#
# Example:
#   powershell -ExecutionPolicy Bypass -File scripts\setup-caption.ps1
#   powershell -File scripts\setup-caption.ps1 -SkipDownload

[CmdletBinding()]
param(
  [string]$ModelsRoot = '',
  [string]$LlamaServer = '',
  [switch]$SkipDownload,
  [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-ModelsRoot([string]$Hint) {
  $shared = Join-Path $env:LOCALAPPDATA 'Comfy-Desktop\ComfyUI-Shared\models'
  $candidates = @(
    $Hint,
    $env:COMFY_MODELS_ROOT,
    $shared
  ) | 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 @{ Name = $name; Source = $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
  }
}

Write-Step 'Resolve models root'
$modelsRoot = Resolve-ModelsRoot $ModelsRoot
$captionDir = Join-Path $modelsRoot 'caption\qwen25vl-7b-nsfw-v4'
Ensure-Dir $captionDir
Write-Ok "Models root: $modelsRoot"
Write-Ok "Caption dir: $captionDir"

$model = 'Qwen2.5-VL-7B-NSFW-Caption-V4.Q5_K_M.gguf'
$mmproj = 'Qwen2.5-VL-7B-NSFW-Caption-V4.mmproj-f16.gguf'
$modelPath = Join-Path $captionDir $model
$mmprojPath = Join-Path $captionDir $mmproj

if (-not $SkipDownload) {
  Write-Step 'Download GGUF + mmproj only (~7GB)'
  $needModel = $Force -or -not (Test-Path $modelPath)
  $needMmproj = $Force -or -not (Test-Path $mmprojPath)
  if (-not $needModel -and -not $needMmproj) {
    Write-Ok 'Both weight files already present.'
  } else {
    $hf = Find-HfCli
    if ($hf.Name -eq 'hf') {
      & $hf.Source download mradermacher/Qwen2.5-VL-7B-NSFW-Caption-V4-GGUF `
        --include $model `
        --include $mmproj `
        --local-dir $captionDir
    } else {
      & $hf.Source download mradermacher/Qwen2.5-VL-7B-NSFW-Caption-V4-GGUF `
        --include $model `
        --include $mmproj `
        --local-dir $captionDir
    }
    if (-not (Test-Path $modelPath)) { throw "Missing $model after download." }
    if (-not (Test-Path $mmprojPath)) { throw "Missing $mmproj after download." }
    Write-Ok 'Download complete.'
  }
} else {
  Write-Warn 'SkipDownload set — not fetching weights.'
}

Write-Step 'Locate llama-server'
$server = $LlamaServer
if (-not $server) { $server = $env:CAPTION_LLAMA_SERVER }
if (-not $server) {
  $cmd = Get-Command llama-server -ErrorAction SilentlyContinue
  if ($cmd) { $server = $cmd.Source }
}
if ($server) {
  Write-Ok "llama-server: $server"
} else {
  Write-Warn 'llama-server not on PATH. Install with: winget install ggml.llamacpp'
  Write-Warn 'Or set CAPTION_LLAMA_SERVER to a CUDA build of llama-server.exe'
}

$ready = @{
  model = $modelPath
  mmproj = $mmprojPath
  modelsDir = $captionDir
  llamaServer = $server
  backend = 'llama.cpp'
  updatedAt = (Get-Date).ToString('o')
}
$readyPath = Join-Path $captionDir 'aigen-ready.json'
$ready | ConvertTo-Json | Set-Content -Path $readyPath -Encoding utf8
Write-Ok "Wrote $readyPath"
Write-Host ""
Write-Host "Caption host looks for:" -ForegroundColor Cyan
Write-Host "  $modelPath"
Write-Host "  $mmprojPath"
Write-Host "Restart the Comfy host agent after install so /caption/status reports configured=true."