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

$ErrorActionPreference = "Stop"

if (-not $LlamaDir) {
    $LlamaDir = Join-Path $env:TEMP "llama.cpp"
}
if (-not $OutDir) {
    $OutDir = Join-Path $env:APPDATA "XianrenStudio\engines"
}

if (-not (Get-Command cmake -ErrorAction SilentlyContinue)) {
    throw "CMake is required. Install it with: winget install Kitware.CMake"
}

if (-not (Test-Path (Join-Path $LlamaDir "CMakeLists.txt"))) {
    Write-Host "Cloning llama.cpp into $LlamaDir"
    git clone --depth 1 https://github.com/ggerganov/llama.cpp.git $LlamaDir
}

function Build-Backend([string]$name, [string]$extraFlag) {
    $buildDir = Join-Path $LlamaDir "build-$name"
    Write-Host "Configuring $name..."
    cmake -S $LlamaDir -B $buildDir -DCMAKE_BUILD_TYPE=Release -DGGML_NATIVE=ON $extraFlag
    Write-Host "Building $name..."
    cmake --build $buildDir --config Release --target llama-server -j 8
    $src = Join-Path $buildDir "bin\Release\llama-server.exe"
    if (-not (Test-Path $src)) { $src = Join-Path $buildDir "bin\llama-server.exe" }
    New-Item -ItemType Directory -Force -Path $OutDir | Out-Null
    Copy-Item $src (Join-Path $OutDir "llama-server-$name.exe") -Force
    Write-Host "Built llama-server-$name.exe"
}

switch ($Backend) {
    "cpu" { Build-Backend "cpu" "-DGGML_CUDA=OFF" }
    "cuda" { Build-Backend "cuda" "-DGGML_CUDA=ON" }
    "vulkan" { Build-Backend "vulkan" "-DGGML_VULKAN=ON" }
    "all" {
        Build-Backend "cpu" "-DGGML_CUDA=OFF"
        if (Get-Command nvcc -ErrorAction SilentlyContinue) { Build-Backend "cuda" "-DGGML_CUDA=ON" }
        if (Get-Command glslc -ErrorAction SilentlyContinue) { Build-Backend "vulkan" "-DGGML_VULKAN=ON" }
    }
}

Write-Host "All requested builds completed in $OutDir"