[CmdletBinding()]
param(
    [string]$InstallRoot = ""
)

$ErrorActionPreference = "Stop"
if (-not $InstallRoot) {
    $InstallRoot = Join-Path $PSScriptRoot "..\.local\llama.cpp"
}
function Get-LlamaVersion([string]$ServerPath) {
    return (& cmd.exe /d /c "`"$ServerPath`" --version 2>&1" | Out-String).Trim()
}
$release = Invoke-RestMethod -Uri "https://api.github.com/repos/ggml-org/llama.cpp/releases/latest"
$tag = [string]$release.tag_name
if (-not $tag) {
    throw "GitHub did not return a latest llama.cpp release tag."
}

$InstallRoot = [System.IO.Path]::GetFullPath($InstallRoot)
$versionDir = Join-Path $InstallRoot $tag
$manifestPath = Join-Path $InstallRoot "installed.json"
if (Test-Path $manifestPath) {
    $installed = Get-Content $manifestPath -Raw | ConvertFrom-Json
    $installedNumber = [int](([string]$installed.tag) -replace "\D", "")
    $latestNumber = [int]($tag -replace "\D", "")
    if ($installedNumber -gt $latestNumber) {
        throw "Refusing to replace newer local llama.cpp $($installed.tag) with $tag."
    }
}

$cudaAssets = @($release.assets | Where-Object {
    $_.name -match "win-cuda-12.*-x64\.zip$"
})
$binaryAsset = $cudaAssets | Where-Object { $_.name -notmatch "cudart" } | Select-Object -First 1
$runtimeAsset = $cudaAssets | Where-Object { $_.name -match "cudart" } | Select-Object -First 1
if (-not $binaryAsset -or -not $runtimeAsset) {
    throw "The latest release does not contain both Windows x64 CUDA 12 binary and runtime archives."
}

New-Item -ItemType Directory -Force -Path $versionDir | Out-Null
$existingServer = Join-Path $versionDir "llama-server.exe"
if (Test-Path $existingServer) {
    $versionText = Get-LlamaVersion $existingServer
    [ordered]@{
        tag = $tag
        installed_at = (Get-Date).ToString("o")
        binary_asset = $binaryAsset.name
        runtime_asset = $runtimeAsset.name
        version = $versionText
    } | ConvertTo-Json | Set-Content -LiteralPath $manifestPath -Encoding UTF8
    $versionText
    exit 0
}

$tempDir = Join-Path ([System.IO.Path]::GetTempPath()) "tower-llama-$tag"
New-Item -ItemType Directory -Force -Path $tempDir | Out-Null
try {
    foreach ($asset in @($binaryAsset, $runtimeAsset)) {
        $archive = Join-Path $tempDir $asset.name
        Invoke-WebRequest -Uri $asset.browser_download_url -OutFile $archive
        Expand-Archive -LiteralPath $archive -DestinationPath $versionDir -Force
    }

    $server = Get-ChildItem -LiteralPath $versionDir -Filter "llama-server.exe" -Recurse |
        Select-Object -First 1
    if (-not $server) {
        throw "llama-server.exe was not found after extraction."
    }
    if ([System.IO.Path]::GetFullPath($server.DirectoryName) -ne [System.IO.Path]::GetFullPath($versionDir)) {
        Get-ChildItem -LiteralPath $server.DirectoryName -Force |
            Copy-Item -Destination $versionDir -Recurse -Force
    }

    $versionText = Get-LlamaVersion (Join-Path $versionDir "llama-server.exe")
    [ordered]@{
        tag = $tag
        installed_at = (Get-Date).ToString("o")
        binary_asset = $binaryAsset.name
        runtime_asset = $runtimeAsset.name
        version = $versionText
    } | ConvertTo-Json | Set-Content -LiteralPath $manifestPath -Encoding UTF8
    $versionText
}
finally {
    if (Test-Path $tempDir) {
        Remove-Item -LiteralPath $tempDir -Recurse -Force
    }
}