param(
    [Parameter(Mandatory = $true)]
    [ValidateSet("start", "stop", "restart", "status", "logs")]
    [string]$Action,

    [string]$DockerHost = "ssh://docker-gpu",
    [string]$SshTarget = "docker-gpu",
    [string]$ContainerName = "llm-llama-qwen3-coder-q6-test",
    [string]$HostPort = "8081",
    [string]$HostModelsDir = "Z:/codex/LLM/models",
    [string]$ModelPath = "/models/gguf/1c/qwen3-coder-30b-a3b-instruct-q6_k/Qwen3-Coder-30B-A3B-Instruct-Q6_K.gguf",
    [string]$ServedModelName = "qwen3-coder-1c-q6",
    [int]$ContextSize = 8192,
    [int]$Threads = 12,
    [int]$Parallel = 1,
    [string]$GpuLayers = "auto",
    [string]$Fit = "on",
    [int]$FitContext = 4096,
    [string]$Reasoning = "auto",
    [string]$ReasoningFormat = "none",
    [string]$LoraPath = "",
    [double]$LoraScale = 1.0,
    [string]$NoMmap = "on",
    [int]$Tail = 80
)

$ErrorActionPreference = "Stop"

function Test-ContainerExists {
    $existing = docker --host $DockerHost ps -a --filter "name=^$ContainerName$" --format "{{.Names}}"
    return $existing -contains $ContainerName
}

function Start-GpuModel {
    if (Test-ContainerExists) {
        docker --host $DockerHost start $ContainerName
        return
    }

    $args = @(
        "--host", $DockerHost,
        "run", "-d",
        "--name", $ContainerName,
        "--restart", "unless-stopped",
        "--gpus", "all",
        "-p", "${HostPort}:8080",
        "-v", "${HostModelsDir}:/models:ro",
        "ghcr.io/ggml-org/llama.cpp:server-cuda",
        "--host", "0.0.0.0",
        "--port", "8080",
        "--model", $ModelPath,
        "--alias", $ServedModelName,
        "--ctx-size", "$ContextSize",
        "--n-gpu-layers", $GpuLayers,
        "--threads", "$Threads",
        "--parallel", "$Parallel",
        "--fit", $Fit,
        "--fit-ctx", "$FitContext",
        "--reasoning", $Reasoning,
        "--reasoning-format", $ReasoningFormat
    )
    if ($NoMmap -in @("on", "true", "1")) {
        $args += "--no-mmap"
    }
    if ($LoraPath) {
        if ($LoraScale -eq 1.0) {
            $args += @("--lora", $LoraPath)
        } else {
            $args += @("--lora-scaled", $LoraPath, "$LoraScale")
        }
    }
    docker @args
}

switch ($Action) {
    "start" {
        Start-GpuModel
    }
    "stop" {
        if (Test-ContainerExists) {
            docker --host $DockerHost stop $ContainerName
        }
    }
    "restart" {
        if (Test-ContainerExists) {
            docker --host $DockerHost restart $ContainerName
        } else {
            Start-GpuModel
        }
    }
    "status" {
        docker --host $DockerHost ps -a --filter "name=^$ContainerName$"
        Write-Host ""
        if ($LoraPath) {
            Write-Host "LoRA: $LoraPath (scale=$LoraScale)"
        } else {
            Write-Host "LoRA: disabled"
        }
        Write-Host ""
        Write-Host "GPU:"
        ssh $SshTarget "nvidia-smi --query-gpu=index,name,memory.used,memory.total,utilization.gpu --format=csv,noheader,nounits"
    }
    "logs" {
        docker --host $DockerHost logs --tail $Tail $ContainerName
    }
}

exit $LASTEXITCODE