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

    [string]$DockerHost = "ssh://test-docker",
    [string]$ContainerName = "llm-qwen3-coder-q6-cpu-test",
    [string]$HostPort = "18086",
    [string]$ModelDir = "/home/test/llm/models/qwen3-coder-30b-a3b-instruct-q6_k",
    [string]$ModelFile = "Qwen3-Coder-30B-A3B-Instruct-Q6_K.gguf",
    [string]$ServedModelName = "qwen3-coder-1c-q6-cpu",
    [int]$Threads = 64,
    [int]$ContextSize = 8192,
    [int]$BatchSize = 512,
    [int]$Tail = 80
)

$ErrorActionPreference = "Stop"

function Start-CpuModel {
    $existing = docker -H $DockerHost ps -a --filter "name=^$ContainerName$" --format "{{.Names}}"
    if ($existing -contains $ContainerName) {
        docker -H $DockerHost start $ContainerName
        return
    }

    $args = @(
        "-H", $DockerHost, "run", "-d",
        "--name", $ContainerName,
        "--restart", "unless-stopped",
        "-p", "${HostPort}:8080",
        "-v", "${ModelDir}:/models:ro",
        "ghcr.io/ggml-org/llama.cpp:server",
        "--host", "0.0.0.0",
        "--port", "8080",
        "--model", "/models/$ModelFile",
        "--alias", $ServedModelName,
        "--ctx-size", "$ContextSize",
        "--threads", "$Threads",
        "--parallel", "1",
        "--batch-size", "$BatchSize"
    )
    docker @args
}

switch ($Action) {
    "start" {
        Start-CpuModel
    }
    "stop" {
        docker -H $DockerHost stop $ContainerName
    }
    "restart" {
        $existing = docker -H $DockerHost ps -a --filter "name=^$ContainerName$" --format "{{.Names}}"
        if ($existing -contains $ContainerName) {
            docker -H $DockerHost restart $ContainerName
        } else {
            Start-CpuModel
        }
    }
    "status" {
        docker -H $DockerHost ps -a --filter "name=^$ContainerName$"
        Write-Host ""
        docker -H $DockerHost stats --no-stream --format "table {{.Name}}`t{{.CPUPerc}}`t{{.MemUsage}}" $ContainerName 2>$null
    }
    "logs" {
        docker -H $DockerHost logs --tail $Tail $ContainerName
    }
}

exit $LASTEXITCODE
