[CmdletBinding()]
param(
    [Parameter()]
    [string]$ModelDir = 'C:\Models\Ornith-MTP',

    [Parameter()]
    [ValidateRange(1, 65535)]
    [int]$Port = 18081,

    [Parameter()]
    [ValidateRange(4096, 131072)]
    [int]$Context = 16384,

    [Parameter()]
    [string]$ContainerName = 'ornith-mtp-api'
)

$ErrorActionPreference = 'Stop'

$Image = 'ghcr.io/ggml-org/llama.cpp@sha256:1b3d1458ccda7287feab41b8001311acc03e24cde99ec0a2908fe83830562f38'
$ModelFile = 'ornith-1.0-35b-MTP-graft-down-Q4_0.gguf'
$ModelPath = Join-Path $ModelDir $ModelFile

if (-not (Get-Command docker -ErrorAction SilentlyContinue)) {
    throw 'docker.exe was not found. Install and start Docker Desktop first.'
}

if (-not (Test-Path -LiteralPath $ModelPath -PathType Leaf)) {
    throw "Model file was not found: $ModelPath"
}

$existing = docker ps -a --filter "name=^/${ContainerName}$" --format '{{.Names}}'
if ($LASTEXITCODE -ne 0) {
    throw 'Docker is unavailable. Start Docker Desktop and verify WSL2/NVIDIA support.'
}
if ($existing -contains $ContainerName) {
    throw "Container '$ContainerName' already exists. Stop it with: docker stop $ContainerName"
}

Write-Host "Starting $ContainerName on http://127.0.0.1:$Port ..."

$DockerArgs = @(
    'run', '--detach', '--rm',
    '--name', $ContainerName,
    '--gpus', 'all',
    '--publish', "127.0.0.1:${Port}:8080",
    '--mount', "type=bind,source=$ModelDir,target=/models,readonly",
    $Image,
    '--model', "/models/$ModelFile",
    '--alias', 'ornith-1.0-35b-mtp',
    '--host', '0.0.0.0',
    '--port', '8080',
    '--ctx-size', [string]$Context,
    '--parallel', '1',
    '--n-gpu-layers', 'all',
    '--cpu-moe',
    '--no-mmap',
    '--batch-size', '2048',
    '--ubatch-size', '512',
    '--flash-attn', 'on',
    '--cache-type-k', 'q8_0',
    '--cache-type-v', 'q8_0',
    '--jinja',
    '--metrics',
    '--spec-type', 'ngram-mod,draft-mtp',
    '--spec-draft-n-max', '2',
    '--spec-draft-n-min', '1',
    '--spec-draft-p-min', '0.20',
    '--spec-ngram-mod-n-min', '48',
    '--spec-ngram-mod-n-max', '64',
    '--spec-ngram-mod-n-match', '24'
)

$ContainerId = & docker @DockerArgs
if ($LASTEXITCODE -ne 0) {
    throw 'docker run failed. Review the Docker Desktop and NVIDIA configuration.'
}
Write-Host "Container: $ContainerId"

$HealthUrl = "http://127.0.0.1:$Port/health"
$Deadline = (Get-Date).AddMinutes(5)
do {
    Start-Sleep -Seconds 2
    try {
        $Health = Invoke-RestMethod -Uri $HealthUrl -TimeoutSec 5
        if ($Health.status -eq 'ok') {
            Write-Host 'Ornith endpoint is ready.' -ForegroundColor Green
            Write-Host "OpenAI base URL: http://127.0.0.1:$Port/v1"
            Write-Host 'Model: ornith-1.0-35b-mtp'
            Write-Host "Logs: docker logs -f $ContainerName"
            Write-Host "Stop: docker stop $ContainerName"
            exit 0
        }
    } catch {
        # Model loading is still in progress.
    }
} while ((Get-Date) -lt $Deadline)

Write-Warning "Endpoint did not become ready within five minutes. Inspect: docker logs $ContainerName"
exit 1