#!/usr/bin/env pwsh
# setup.ps1 — one-time install of ComfyUI + kohya_ss

$root = $PSScriptRoot
Set-Location $root

Write-Host "==> setting up animepics environment" -ForegroundColor Cyan

# create model dirs
$model_dirs = @(
    "models/checkpoints",
    "models/loras",
    "models/vae",
    "models/embeddings",
    "models/controlnet",
    "models/upscale_models",
    "training_data",
    "output"
)
foreach ($dir in $model_dirs) {
    New-Item -ItemType Directory -Force -Path "$root/$dir" | Out-Null
}
Write-Host "  created model directories" -ForegroundColor Green

# --- ComfyUI ---
Write-Host "`n==> installing ComfyUI" -ForegroundColor Cyan
if (-not (Test-Path "$root/comfyui")) {
    git clone https://github.com/comfyanonymous/ComfyUI.git "$root/comfyui"
} else {
    Write-Host "  comfyui already cloned, pulling latest" -ForegroundColor Yellow
    git -C "$root/comfyui" pull
}

Push-Location "$root/comfyui"
python -m venv venv
.\venv\Scripts\Activate.ps1
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124
pip install -r requirements.txt
deactivate
Pop-Location
Write-Host "  ComfyUI installed" -ForegroundColor Green

# ComfyUI custom nodes
Write-Host "`n==> installing ComfyUI custom nodes" -ForegroundColor Cyan
$custom_nodes = @(
    @{ name="ComfyUI-Manager";      url="https://github.com/ltdrdata/ComfyUI-Manager.git" },
    @{ name="ComfyUI_IPAdapter_plus"; url="https://github.com/cubiq/ComfyUI_IPAdapter_plus.git" },
    @{ name="ComfyUI-Advanced-ControlNet"; url="https://github.com/Kosinkadink/ComfyUI-Advanced-ControlNet.git" },
    @{ name="was-node-suite-comfyui"; url="https://github.com/WASasquatch/was-node-suite-comfyui.git" }
)
New-Item -ItemType Directory -Force -Path "$root/comfyui/custom_nodes" | Out-Null
foreach ($node in $custom_nodes) {
    $node_path = "$root/comfyui/custom_nodes/$($node.name)"
    if (-not (Test-Path $node_path)) {
        git clone $node.url $node_path
        Write-Host "  cloned $($node.name)" -ForegroundColor Green
    } else {
        git -C $node_path pull
        Write-Host "  updated $($node.name)" -ForegroundColor Yellow
    }
}

# copy extra_model_paths config so ComfyUI finds our shared models dir
Copy-Item -Force "$root/config/extra_model_paths.yaml" "$root/comfyui/extra_model_paths.yaml"
Write-Host "  copied extra_model_paths.yaml" -ForegroundColor Green

# --- kohya_ss ---
Write-Host "`n==> installing kohya_ss" -ForegroundColor Cyan
if (-not (Test-Path "$root/kohya_ss")) {
    git clone https://github.com/bmaltais/kohya_ss.git "$root/kohya_ss"
} else {
    Write-Host "  kohya_ss already cloned, pulling latest" -ForegroundColor Yellow
    git -C "$root/kohya_ss" pull
}

Push-Location "$root/kohya_ss"
python -m venv venv
.\venv\Scripts\Activate.ps1
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124
pip install -r requirements.txt
deactivate
Pop-Location
Write-Host "  kohya_ss installed" -ForegroundColor Green

Write-Host "`n==> setup complete!" -ForegroundColor Cyan
Write-Host "  next: download your base model into models/checkpoints/" -ForegroundColor White
Write-Host "  recommended: NoobAI-XL from https://civitai.com/models/833294" -ForegroundColor White
Write-Host "`n  then run: .\launch_comfyui.ps1" -ForegroundColor White