param(
    [string]$SdScriptsRepo = "https://github.com/kohya-ss/sd-scripts.git",
    [string]$BackendDir = "",
    [string]$Python = "python",
    [switch]$SkipTorch
)

$ErrorActionPreference = "Stop"

$PluginDir = Split-Path -Parent (Split-Path -Parent $MyInvocation.MyCommand.Path)
if ([string]::IsNullOrWhiteSpace($BackendDir)) {
    $BackendDir = Join-Path $PluginDir "backend\sd-scripts"
}

$BackendDir = [System.IO.Path]::GetFullPath($BackendDir)
$VenvDir = Join-Path $BackendDir ".venv"
$VenvPython = Join-Path $VenvDir "Scripts\python.exe"

Write-Host "AnimaForge Windows backend setup"
Write-Host "Backend: $BackendDir"

if (-not (Get-Command git -ErrorAction SilentlyContinue)) {
    throw "git was not found. Install Git for Windows first."
}

if (-not (Test-Path $BackendDir)) {
    New-Item -ItemType Directory -Force -Path (Split-Path -Parent $BackendDir) | Out-Null
    git clone $SdScriptsRepo $BackendDir
} else {
    Write-Host "sd-scripts backend already exists. Pulling latest changes..."
    git -C $BackendDir pull --ff-only
}

if (-not (Test-Path $VenvPython)) {
    Write-Host "Creating backend virtual environment..."
    & $Python -m venv $VenvDir
}

Write-Host "Upgrading pip..."
& $VenvPython -m pip install --upgrade pip setuptools wheel

if (-not $SkipTorch) {
    Write-Host "Installing PyTorch CUDA 12.1 wheels. Edit this script if your CUDA/PyTorch requirement is different."
    & $VenvPython -m pip install --index-url https://download.pytorch.org/whl/cu121 torch torchvision
}

Write-Host "Installing sd-scripts requirements..."
& $VenvPython -m pip install -r (Join-Path $BackendDir "requirements.txt")
& $VenvPython -m pip install -e $BackendDir

Write-Host ""
Write-Host "Backend setup complete."
Write-Host "Training Python: $VenvPython"
Write-Host "Set ANIMAFORGE_BACKEND_DIR only if you move the backend somewhere else."