[CmdletBinding()]
param(
    [ValidateSet("vision-detector", "vision-segmenter", "change-detector", "thermal-analyzer")]
    [string]$ModelGroup,
    [string]$ModelPath,
    [string]$SecondaryModelPath,
    [string]$LabelsPath,
    [switch]$Restart
)

$ErrorActionPreference = "Stop"
$root = Split-Path -Parent $PSScriptRoot
$registryPath = Join-Path $root "infra\model-registry\cpu-models.json"
$modelRoot = Join-Path $root "runtime\models"
$registry = Get-Content $registryPath -Raw -Encoding UTF8 | ConvertFrom-Json

New-Item -ItemType Directory -Force -Path $modelRoot | Out-Null

$rows = foreach ($model in $registry.models) {
    $secondary = $null
    if ($model.artifact) {
        $artifact = Join-Path $modelRoot ([string]$model.artifact).Replace("/", "\")
        New-Item -ItemType Directory -Force -Path (Split-Path -Parent $artifact) | Out-Null
    } else {
        $artifact = $null
    }
    if ($model.secondary_artifact) {
        $secondary = Join-Path $modelRoot ([string]$model.secondary_artifact).Replace("/", "\")
        New-Item -ItemType Directory -Force -Path (Split-Path -Parent $secondary) | Out-Null
    }
    if ($model.labels) {
        $labels = Join-Path $modelRoot ([string]$model.labels).Replace("/", "\")
        New-Item -ItemType Directory -Force -Path (Split-Path -Parent $labels) | Out-Null
    } else {
        $labels = $null
    }
    [PSCustomObject]@{
        ModelGroup = $model.model_group
        Model = $model.display_name
        Runtime = $model.runtime
        Artifact = $artifact
        SecondaryArtifact = $secondary
        Installed = $(if ($artifact) { Test-Path -LiteralPath $artifact -PathType Leaf } else { $true }) -and
            $(if ($secondary) { Test-Path -LiteralPath $secondary -PathType Leaf } else { $true })
        Labels = $labels
    }
}

if ($ModelGroup) {
    if (-not $ModelPath) {
        throw "-ModelPath is required when -ModelGroup is specified."
    }
    $source = (Resolve-Path -LiteralPath $ModelPath).Path
    if ([IO.Path]::GetExtension($source) -ne ".onnx") {
        throw "The CPU model artifact must be an .onnx file."
    }
    $target = $rows | Where-Object { $_.ModelGroup -eq $ModelGroup }
    if (-not $target -or -not $target.Artifact) {
        throw "Model group $ModelGroup has no ONNX artifact location."
    }
    Copy-Item -LiteralPath $source -Destination $target.Artifact -Force
    if ($SecondaryModelPath) {
        if (-not $target.SecondaryArtifact) {
            throw "Model group $ModelGroup does not define a secondary ONNX artifact."
        }
        $secondarySource = (Resolve-Path -LiteralPath $SecondaryModelPath).Path
        if ([IO.Path]::GetExtension($secondarySource) -ne ".onnx") {
            throw "The secondary CPU model artifact must be an .onnx file."
        }
        Copy-Item -LiteralPath $secondarySource -Destination $target.SecondaryArtifact -Force
        Write-Host "Installed $ModelGroup secondary artifact: $($target.SecondaryArtifact)" -ForegroundColor Green
    } elseif ($target.SecondaryArtifact -and -not (Test-Path -LiteralPath $target.SecondaryArtifact -PathType Leaf)) {
        Write-Warning "$ModelGroup also requires -SecondaryModelPath: $($target.SecondaryArtifact)"
    }
    if ($LabelsPath) {
        if (-not $target.Labels) {
            throw "Model group $ModelGroup does not define a label file location."
        }
        $labelSource = (Resolve-Path -LiteralPath $LabelsPath).Path
        Copy-Item -LiteralPath $labelSource -Destination $target.Labels -Force
    }
    $target.Installed = (Test-Path -LiteralPath $target.Artifact -PathType Leaf) -and
        (-not $target.SecondaryArtifact -or (Test-Path -LiteralPath $target.SecondaryArtifact -PathType Leaf))
    Write-Host "Installed $ModelGroup artifact: $($target.Artifact)" -ForegroundColor Green
}

$rows | Select-Object ModelGroup, Model, Runtime, Installed, Artifact, SecondaryArtifact | Format-Table -AutoSize

$layoutPath = Join-Path $modelRoot "MODEL_LAYOUT.txt"
$layout = @(
    "Rail UAV Inspection CPU model directory",
    "Profile: $($registry.profile)",
    "Generated: $([DateTimeOffset]::Now.ToString('o'))",
    "",
    "The labels.txt line order must exactly match the exported model class order.",
    "Production railway weights are not downloaded automatically."
)
Set-Content -LiteralPath $layoutPath -Value $layout -Encoding UTF8

if ($Restart) {
    docker compose -f "$root\infra\docker-compose.yml" -f "$root\infra\docker-compose.cpu.yml" restart vision-inference pointcloud-analysis
    if ($LASTEXITCODE -ne 0) {
        throw "Failed to restart AI services."
    }
}
