[CmdletBinding()]
param(
    [Parameter(Mandatory = $true)]
    [string]$JobRoot,

    [Parameter(Mandatory = $true)]
    [string]$RunnerPath,

    [string]$RuntimeRoot = "D:\NDC_MISSIONCORE\runtime",

    [string]$CalibrationSha256 = "05f3ad9b38b3a4fc95388a8ec83da83c745e217709e51787b3d5aad0969f6fa9",

    [string]$CalibrationSlot = "camera_1"
)

$ErrorActionPreference = "Stop"
$ProgressPreference = "SilentlyContinue"

function Assert-LastExitCode {
    param([string]$Operation)
    if ($LASTEXITCODE -ne 0) {
        throw "$Operation failed with exit code $LASTEXITCODE"
    }
}

function Assert-RegularFile {
    param([string]$Path, [string]$Label)
    $item = Get-Item -LiteralPath $Path -Force
    if (-not $item.PSIsContainer -and -not ($item.Attributes -band [IO.FileAttributes]::ReparsePoint)) {
        return $item.FullName
    }
    throw "$Label is not a regular file"
}

function Assert-Directory {
    param([string]$Path, [string]$Label)
    $item = Get-Item -LiteralPath $Path -Force
    if ($item.PSIsContainer -and -not ($item.Attributes -band [IO.FileAttributes]::ReparsePoint)) {
        return $item.FullName
    }
    throw "$Label is not a real directory"
}

function Convert-ToDockerPath {
    param([string]$Path)
    return $Path.Replace("\", "/")
}

$jobDirectory = Assert-Directory (Resolve-Path -LiteralPath $JobRoot).Path "Job root"
$runner = Assert-RegularFile (Resolve-Path -LiteralPath $RunnerPath).Path "Runner"
$runtime = Assert-Directory (Resolve-Path -LiteralPath $RuntimeRoot).Path "Runtime root"
$job = Get-Content -LiteralPath (Join-Path $jobDirectory "job.json") -Raw | ConvertFrom-Json
if ($job.schema_version -ne "missioncore.compute-job/v1" -or $job.job_id -ne (Split-Path $jobDirectory -Leaf)) {
    throw "Compute job manifest is incompatible"
}
if ($CalibrationSha256 -notmatch "^[a-f0-9]{64}$" -or $CalibrationSlot -notmatch "^[A-Za-z0-9._-]+$") {
    throw "Calibration binding is invalid"
}
Write-Output "PHASE=job-manifest-validated"

$sourceId = [string]$job.input.source_id
$epoch = [int]$job.input.codec_epoch
$frameCount = [int]$job.input.segment_count
$timelineStart = [double]$job.input.timeline.start_seconds
$timelineEnd = [double]$job.input.timeline.end_seconds
$timelineDuration = $timelineEnd - $timelineStart
if ($frameCount -lt 1 -or $timelineDuration -le 0) {
    throw "Compute job frame/timeline contract is invalid"
}

$epochRoot = Join-Path $jobDirectory ("input\camera\{0}\epoch-{1}" -f $sourceId, $epoch)
$epochRoot = Assert-Directory $epochRoot "Camera epoch"
$initPath = Assert-RegularFile (Join-Path $epochRoot "init.mp4") "Camera init"
$segmentsRoot = Assert-Directory (Join-Path $epochRoot "segments") "Camera segments"
$derivedRoot = Join-Path $runtime "derived"
$tmpRoot = Join-Path $runtime "tmp"
$cacheRoot = Assert-Directory (Join-Path $runtime "cache\perception-p0-models-v1") "Model cache"
$torchEnvironment = Assert-Directory (Join-Path $derivedRoot "perception-p0-env-v1") "Torch environment"
$transformersEnvironment = Assert-Directory (Join-Path $derivedRoot "perception-p0-transformers4576-v1") "Transformers environment"
$runnerRoot = Split-Path $runner -Parent
$runnerName = Split-Path $runner -Leaf
$null = New-Item -ItemType Directory -Path $derivedRoot -Force
$null = New-Item -ItemType Directory -Path $tmpRoot -Force

$preflightArgs = @(
    "run", "--rm", "--gpus", "all", "--network", "none",
    "--security-opt", "no-new-privileges:true", "--cap-drop", "ALL",
    "-e", "PYTHONPATH=/opt/transformers:/opt/env",
    "-e", "TORCH_HOME=/cache/torch",
    "-e", "HF_HOME=/cache/huggingface",
    "-e", "HF_HUB_OFFLINE=1",
    "-e", "TRANSFORMERS_OFFLINE=1",
    "-v", ((Convert-ToDockerPath $jobDirectory) + ":/job:ro"),
    "-v", ((Convert-ToDockerPath $cacheRoot) + ":/cache:rw"),
    "-v", ((Convert-ToDockerPath $torchEnvironment) + ":/opt/env:ro"),
    "-v", ((Convert-ToDockerPath $transformersEnvironment) + ":/opt/transformers:ro"),
    "-v", ((Convert-ToDockerPath $runnerRoot) + ":/runner:ro"),
    "--entrypoint", "python3",
    "nvcr.io/nvidia/tritonserver:26.06-py3",
    ("/runner/{0}" -f $runnerName), "preflight",
    "--job", "/job",
    "--cache", "/cache"
)
Write-Output "PHASE=worker-preflight-started"
& docker @preflightArgs
Assert-LastExitCode "Worker preflight"
Write-Output "PHASE=worker-preflight-complete"

$runToken = [Guid]::NewGuid().ToString("N")
$workRoot = Join-Path $tmpRoot ("{0}-panoptic-{1}" -f $job.job_id, $runToken)
$framesRoot = Join-Path $workRoot "frames"
$streamPath = Join-Path $workRoot "camera.mp4"
$ptsPath = Join-Path $workRoot "pts.json"
$timelinePath = Join-Path $workRoot "timeline.jsonl"
$publishRoot = Join-Path $derivedRoot (".{0}-panoptic-{1}.publish" -f $job.job_id, $runToken)
$stagingRoot = Join-Path $publishRoot "output"
$null = New-Item -ItemType Directory -Path $framesRoot
$null = New-Item -ItemType Directory -Path $publishRoot
$totalWatch = [Diagnostics.Stopwatch]::StartNew()
Write-Output "PHASE=private-staging-created"

try {
    $stream = [IO.File]::Open($streamPath, [IO.FileMode]::CreateNew, [IO.FileAccess]::Write, [IO.FileShare]::None)
    try {
        foreach ($path in @($initPath)) {
            $input = [IO.File]::OpenRead($path)
            try {
                $input.CopyTo($stream)
            }
            finally {
                $input.Dispose()
            }
        }
        for ($sequence = 1; $sequence -le $frameCount; $sequence++) {
            $path = Assert-RegularFile (Join-Path $segmentsRoot ("{0}.m4s" -f $sequence)) "Camera segment"
            $input = [IO.File]::OpenRead($path)
            try {
                $input.CopyTo($stream)
            }
            finally {
                $input.Dispose()
            }
            if ($sequence % 500 -eq 0 -or $sequence -eq $frameCount) {
                Write-Output ("PHASE=stream-reconstruction SEGMENTS={0}/{1}" -f $sequence, $frameCount)
            }
        }
        $stream.Flush($true)
    }
    finally {
        $stream.Dispose()
    }

    $extractWatch = [Diagnostics.Stopwatch]::StartNew()
    Write-Output "PHASE=frame-extraction-started"
    # This archive contains valid strictly-increasing PTS but image2 rounds DTS
    # to its own time base and reports harmless duplicate-DTS diagnostics. Exact
    # frame-count and JSON timestamp checks below are the admission boundary.
    & ffmpeg -hide_banner -loglevel fatal -i $streamPath -map 0:v:0 -fps_mode passthrough (Join-Path $framesRoot "frame-%06d.png")
    Assert-LastExitCode "Camera extraction"
    & ffprobe -v error -select_streams v:0 -show_entries frame=best_effort_timestamp_time -of json $streamPath | Set-Content -LiteralPath $ptsPath -Encoding utf8
    Assert-LastExitCode "Camera timestamp probe"

    $decodedFrames = @(Get-ChildItem -LiteralPath $framesRoot -File -Filter "frame-*.png")
    $ptsDocument = Get-Content -LiteralPath $ptsPath -Raw | ConvertFrom-Json
    $pts = @($ptsDocument.frames)
    if ($decodedFrames.Count -ne $frameCount -or $pts.Count -ne $frameCount) {
        throw "Decoded frame count differs from the compute job"
    }
    $firstEpochSeconds = [double]::Parse(
        ([string]$pts[0].best_effort_timestamp_time).Trim(),
        [Globalization.CultureInfo]::InvariantCulture
    )
    $previousEpochSeconds = -1.0
    $timelineWriter = [IO.StreamWriter]::new($timelinePath, $false, [Text.UTF8Encoding]::new($false))
    try {
        for ($index = 0; $index -lt $frameCount; $index++) {
            $epochSeconds = [double]::Parse(
                ([string]$pts[$index].best_effort_timestamp_time).Trim(),
                [Globalization.CultureInfo]::InvariantCulture
            ) - $firstEpochSeconds
            if ($epochSeconds -le $previousEpochSeconds -or $epochSeconds -gt ($timelineDuration + 0.001)) {
                throw "Decoded camera timestamps are not strictly monotonic inside the compute job timeline"
            }
            $sessionSeconds = $timelineStart + $epochSeconds
            $row = [ordered]@{
                frame_index = $index
                epoch_seconds = $epochSeconds
                session_seconds = $sessionSeconds
            }
            $timelineWriter.WriteLine(($row | ConvertTo-Json -Compress))
            $previousEpochSeconds = $epochSeconds
        }
        $timelineWriter.Flush()
    }
    finally {
        $timelineWriter.Dispose()
    }
    $extractWatch.Stop()
    Write-Output ("PHASE=frame-extraction-complete FRAMES={0}" -f $frameCount)

    $dockerArgs = @(
        "run", "--rm", "--gpus", "all", "--network", "none",
        "--security-opt", "no-new-privileges:true", "--cap-drop", "ALL",
        "--shm-size", "2g",
        "-e", "PYTHONPATH=/opt/transformers:/opt/env",
        "-e", "TORCH_HOME=/cache/torch",
        "-e", "HF_HOME=/cache/huggingface",
        "-e", "HF_HUB_OFFLINE=1",
        "-e", "TRANSFORMERS_OFFLINE=1",
        "-v", ((Convert-ToDockerPath $jobDirectory) + ":/job:ro"),
        "-v", ((Convert-ToDockerPath $framesRoot) + ":/frames:ro"),
        "-v", ((Convert-ToDockerPath $workRoot) + ":/work:ro"),
        "-v", ((Convert-ToDockerPath $publishRoot) + ":/publish:rw"),
        "-v", ((Convert-ToDockerPath $cacheRoot) + ":/cache:rw"),
        "-v", ((Convert-ToDockerPath $torchEnvironment) + ":/opt/env:ro"),
        "-v", ((Convert-ToDockerPath $transformersEnvironment) + ":/opt/transformers:ro"),
        "-v", ((Convert-ToDockerPath $runnerRoot) + ":/runner:ro"),
        "--entrypoint", "python3",
        "nvcr.io/nvidia/tritonserver:26.06-py3",
        ("/runner/{0}" -f $runnerName), "run",
        "--job", "/job",
        "--frames", "/frames",
        "--timeline", "/work/timeline.jsonl",
        "--output", "/publish/output",
        "--cache", "/cache",
        "--calibration-sha256", $CalibrationSha256,
        "--calibration-slot", $CalibrationSlot,
        "--telemetry-interval-seconds", "1"
    )
    & docker @dockerArgs
    Assert-LastExitCode "Panoptic inference"
    Write-Output "PHASE=panoptic-inference-complete"

    $videoPath = Join-Path $stagingRoot "perception.mp4"
    $encodeWatch = [Diagnostics.Stopwatch]::StartNew()
    $averageFps = $frameCount / $timelineDuration
    $fpsText = $averageFps.ToString("0.#########", [Globalization.CultureInfo]::InvariantCulture)
    $durationText = $timelineDuration.ToString("0.#########", [Globalization.CultureInfo]::InvariantCulture)
    & ffmpeg -hide_banner -loglevel error -framerate $fpsText -i (Join-Path $stagingRoot "overlay-frames\frame-%06d.png") -t $durationText -c:v h264_nvenc -preset p4 -tune hq -rc vbr -cq 21 -b:v 0 -pix_fmt yuv420p -movflags +faststart $videoPath
    Assert-LastExitCode "Panoptic video encoding"
    $encodeWatch.Stop()
    Write-Output "PHASE=video-encoding-complete"

    $masksPath = Join-Path $stagingRoot "masks.tar.gz"
    & tar.exe -czf $masksPath -C $stagingRoot instance-masks semantic-masks
    Assert-LastExitCode "Mask archive publication"

    $totalWatch.Stop()
    $finalizeArgs = @(
        "run", "--rm", "--network", "none",
        "--security-opt", "no-new-privileges:true", "--cap-drop", "ALL",
        "-v", ((Convert-ToDockerPath $stagingRoot) + ":/output:rw"),
        "-v", ((Convert-ToDockerPath $runnerRoot) + ":/runner:ro"),
        "--entrypoint", "python3",
        "nvcr.io/nvidia/tritonserver:26.06-py3",
        ("/runner/{0}" -f $runnerName), "finalize",
        "--output", "/output",
        "--video", "/output/perception.mp4",
        "--masks", "/output/masks.tar.gz",
        "--extract-seconds", $extractWatch.Elapsed.TotalSeconds.ToString("0.######", [Globalization.CultureInfo]::InvariantCulture),
        "--encode-seconds", $encodeWatch.Elapsed.TotalSeconds.ToString("0.######", [Globalization.CultureInfo]::InvariantCulture),
        "--wall-seconds", $totalWatch.Elapsed.TotalSeconds.ToString("0.######", [Globalization.CultureInfo]::InvariantCulture),
        "--encoder", "ffmpeg-h264_nvenc-p4-cq21-yuv420p-faststart"
    )
    & docker @finalizeArgs
    Assert-LastExitCode "Result finalization"

    $result = Get-Content -LiteralPath (Join-Path $stagingRoot "result.json") -Raw | ConvertFrom-Json
    if ($result.schema_version -ne "missioncore.recorded-perception-result/v2" -or $result.result_id -notmatch "^result-[a-f0-9]{64}$") {
        throw "Final result manifest is incompatible"
    }
    $finalRoot = Join-Path $derivedRoot ([string]$result.result_id)
    if (Test-Path -LiteralPath $finalRoot) {
        throw "An immutable result with the same identity already exists: $finalRoot"
    }

    foreach ($temporaryChild in @("overlay-frames", "instance-masks", "semantic-masks")) {
        $temporaryPath = Join-Path $stagingRoot $temporaryChild
        if (Test-Path -LiteralPath $temporaryPath) {
            Remove-Item -LiteralPath $temporaryPath -Recurse -Force
        }
    }
    Move-Item -LiteralPath $stagingRoot -Destination $finalRoot
    Remove-Item -LiteralPath $publishRoot -Force
    Write-Output ("RESULT_ROOT={0}" -f $finalRoot)
    Write-Output ("RESULT_ID={0}" -f $result.result_id)
}
finally {
    if (Test-Path -LiteralPath $workRoot) {
        Remove-Item -LiteralPath $workRoot -Recurse -Force
    }
}
