param(
    [ValidateSet("cpu", "cuda", "vulkan")]
    [string]$Backend = "cpu",
    [string]$Tag = "",
    [string]$OutDir = ""
)

$ErrorActionPreference = "Stop"

if (-not $OutDir) {
    $OutDir = Join-Path $env:APPDATA "XianrenStudio\engines"
}
New-Item -ItemType Directory -Force -Path $OutDir | Out-Null

$repo = "ggerganov/llama.cpp"
$api = "https://api.github.com/repos/$repo/releases/latest"

if (-not $Tag) {
    Write-Host "Querying latest release from GitHub..."
    $release = Invoke-RestMethod -Uri $api -Headers @{ "User-Agent" = "xianren-studio" }
    $Tag = $release.tag_name
    Write-Host "Latest tag: $Tag"
}

$assetName = "llama-$Tag-bin-win-$Backend-x64.zip"
$assetUrl = "https://github.com/$repo/releases/download/$Tag/$assetName"
$zipPath = Join-Path $env:TEMP $assetName

Write-Host "Downloading $assetUrl"
curl.exe -L --fail --progress-bar -o $zipPath $assetUrl
if ($LASTEXITCODE -ne 0) {
    throw "download failed with exit code $LASTEXITCODE"
}

Write-Host "Extracting llama-server.exe to $OutDir"
$extract = Join-Path $env:TEMP "llama-$Backend-extract"
if (Test-Path $extract) { Remove-Item -Recurse -Force $extract }
Expand-Archive -Path $zipPath -DestinationPath $extract

$server = Get-ChildItem -Path $extract -Recurse -Filter "llama-server.exe" | Select-Object -First 1
if (-not $server) {
    throw "llama-server.exe not found in archive"
}
$targetDir = Join-Path $OutDir $Backend
New-Item -ItemType Directory -Force -Path $targetDir | Out-Null
Copy-Item -Path (Join-Path $extract "*") -Destination $targetDir -Recurse -Force
Remove-Item -Recurse -Force $extract

Write-Host "Done. Engine: $(Join-Path $targetDir "llama-server.exe")"