param(
    [Parameter(Mandatory = $true)]
    [ValidateSet("vision-detector", "vision-segmenter", "change-detector", "pointcloud-analyzer", "thermal-analyzer")]
    [string]$ModelGroup,
    [Parameter(Mandatory = $true)]
    [ValidateSet("v1.0.0", "v1.1.0")]
    [string]$ModelVersion,
    [string]$ModelPath,
    [string]$LabelsPath,
    [string]$SecondaryModelPath,
    [switch]$Restart
)

$ErrorActionPreference = "Stop"
$root = Split-Path -Parent $PSScriptRoot
$registryPath = Join-Path $root "infra/model-registry/server-models.json"
$modelRoot = Join-Path $root "runtime/server-models"
$registry = Get-Content -Raw -Encoding UTF8 $registryPath | ConvertFrom-Json
$entry = $registry.models | Where-Object { $_.model_group -eq $ModelGroup -and $_.model_version -eq $ModelVersion } | Select-Object -First 1

if (-not $entry) {
    throw "The model group and version are not registered in server-models.json."
}

if ($ModelGroup -eq "pointcloud-analyzer") {
    Write-Host "Open3D geometry analysis does not require a model artifact."
} else {
    if (-not $ModelPath -or -not (Test-Path -LiteralPath $ModelPath -PathType Leaf)) {
        throw "ModelPath must reference an existing ONNX model file."
    }
    $target = Join-Path $modelRoot ([string]$entry.artifact)
    $targetDirectory = Split-Path -Parent $target
    New-Item -ItemType Directory -Force -Path $targetDirectory | Out-Null
    Copy-Item -LiteralPath $ModelPath -Destination $target -Force
    Write-Host "Installed model: $target"
    Write-Host "SHA256: $((Get-FileHash -LiteralPath $target -Algorithm SHA256).Hash.ToLowerInvariant())"

    if ($entry.labels) {
        if (-not $LabelsPath -or -not (Test-Path -LiteralPath $LabelsPath -PathType Leaf)) {
            throw "LabelsPath is required for this model and must reference an existing file."
        }
        $labelsTarget = Join-Path $modelRoot ([string]$entry.labels)
        New-Item -ItemType Directory -Force -Path (Split-Path -Parent $labelsTarget) | Out-Null
        Copy-Item -LiteralPath $LabelsPath -Destination $labelsTarget -Force
        Write-Host "Installed labels: $labelsTarget"
    }

    if ($entry.secondary_artifact) {
        if (-not $SecondaryModelPath -or -not (Test-Path -LiteralPath $SecondaryModelPath -PathType Leaf)) {
            throw "SecondaryModelPath is required for the thermal analyzer."
        }
        $secondaryTarget = Join-Path $modelRoot ([string]$entry.secondary_artifact)
        New-Item -ItemType Directory -Force -Path (Split-Path -Parent $secondaryTarget) | Out-Null
        Copy-Item -LiteralPath $SecondaryModelPath -Destination $secondaryTarget -Force
        Write-Host "Installed secondary model: $secondaryTarget"
    }
}

if ($Restart) {
    & docker compose --env-file (Join-Path $root "infra/.env.server") `
        -f (Join-Path $root "infra/docker-compose.yml") `
        -f (Join-Path $root "infra/docker-compose.server.yml") `
        restart vision-inference pointcloud-analysis
    if ($LASTEXITCODE -ne 0) { throw "Failed to restart AI services." }
}
