#!/usr/bin/env python3
"""Run and seal a complete recorded panoptic-perception camera epoch."""

from __future__ import annotations

import argparse
import hashlib
import importlib.metadata
import json
import os
import platform
import resource
import statistics
import subprocess
import threading
import time
from datetime import UTC, datetime
from pathlib import Path, PurePosixPath
from typing import Any, TextIO

INSTANCE_SCORE_THRESHOLD = 0.5
INSTANCE_MASK_THRESHOLD = 0.5
SEMANTIC_ALPHA = 0.46
INSTANCE_ALPHA = 0.58
SEMANTIC_MODEL_ID = "microsoft/beit-base-finetuned-ade-640-640"
SEMANTIC_REVISION = "a8b6f5ef4acb2ea55d882989deaa02d39401e2b2"
RESULT_SCHEMA = "missioncore.recorded-perception-result/v2"
REPORT_SCHEMA = "missioncore.perception-run-report/v1"
FRAME_SCHEMA = "missioncore.panoptic-frame/v1"
SAFE_SHA256 = set("0123456789abcdef")


def _arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="command", required=True)
    preflight = subparsers.add_parser("preflight")
    preflight.add_argument("--job", type=Path, required=True)
    preflight.add_argument("--cache", type=Path, required=True)

    run = subparsers.add_parser("run")
    run.add_argument("--job", type=Path, required=True)
    run.add_argument("--frames", type=Path, required=True)
    run.add_argument("--timeline", type=Path, required=True)
    run.add_argument("--output", type=Path, required=True)
    run.add_argument("--cache", type=Path, required=True)
    run.add_argument("--calibration-sha256", required=True)
    run.add_argument("--calibration-slot", required=True)
    run.add_argument("--telemetry-interval-seconds", type=float, default=1.0)

    finalize = subparsers.add_parser("finalize")
    finalize.add_argument("--output", type=Path, required=True)
    finalize.add_argument("--video", type=Path, required=True)
    finalize.add_argument("--masks", type=Path, required=True)
    finalize.add_argument("--extract-seconds", type=float, required=True)
    finalize.add_argument("--encode-seconds", type=float, required=True)
    finalize.add_argument("--wall-seconds", type=float, required=True)
    finalize.add_argument("--encoder", required=True)
    return parser.parse_args()


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _canonical_json(value: object) -> bytes:
    return json.dumps(
        value,
        ensure_ascii=False,
        sort_keys=True,
        separators=(",", ":"),
        allow_nan=False,
    ).encode("utf-8")


def _read_object(path: Path) -> dict[str, Any]:
    value = json.loads(path.read_text(encoding="utf-8-sig"))
    if not isinstance(value, dict):
        raise RuntimeError(f"{path.name} is not a JSON object")
    return value


def _write_json(path: Path, value: object) -> None:
    temporary = path.with_name(f".{path.name}.{os.getpid()}.tmp")
    with temporary.open("x", encoding="utf-8", newline="\n") as stream:
        json.dump(value, stream, ensure_ascii=False, sort_keys=True, indent=2, allow_nan=False)
        stream.write("\n")
        stream.flush()
        os.fsync(stream.fileno())
    os.replace(temporary, path)


def _valid_sha256(value: object) -> bool:
    return isinstance(value, str) and len(value) == 64 and set(value) <= SAFE_SHA256


def _safe_job_path(root: Path, encoded: object) -> Path:
    if not isinstance(encoded, str):
        raise RuntimeError("job artifact path is not a string")
    relative = PurePosixPath(encoded)
    if relative.is_absolute() or ".." in relative.parts or not relative.parts:
        raise RuntimeError("job artifact path is unsafe")
    path = root.joinpath(*relative.parts).resolve(strict=True)
    if not path.is_file() or not path.is_relative_to(root):
        raise RuntimeError("job artifact escapes the job root")
    return path


def _validate_job(job_root: Path) -> dict[str, Any]:
    root = job_root.resolve(strict=True)
    job = _read_object(root / "job.json")
    input_document = job.get("input")
    if (
        job.get("schema_version") != "missioncore.compute-job/v1"
        or not isinstance(input_document, dict)
        or hashlib.sha256(_canonical_json(input_document)).hexdigest() != job.get("input_sha256")
        or job.get("job_id") != f"recorded-camera-{str(job.get('input_sha256'))[:24]}"
    ):
        raise RuntimeError("compute job identity is invalid")
    files = input_document.get("files")
    if not isinstance(files, list) or not files:
        raise RuntimeError("compute job has no files")
    seen: set[str] = set()
    total_bytes = 0
    for artifact in files:
        if not isinstance(artifact, dict):
            raise RuntimeError("compute job artifact descriptor is invalid")
        encoded = artifact.get("path")
        if not isinstance(encoded, str) or encoded in seen:
            raise RuntimeError("compute job artifact descriptor is duplicated")
        seen.add(encoded)
        path = _safe_job_path(root, encoded)
        byte_length = artifact.get("byte_length")
        digest = artifact.get("sha256")
        if (
            not isinstance(byte_length, int)
            or byte_length < 1
            or path.stat().st_size != byte_length
            or not _valid_sha256(digest)
            or _sha256(path) != digest
        ):
            raise RuntimeError(f"compute job artifact changed: {encoded}")
        total_bytes += byte_length
    if total_bytes != input_document.get("byte_length"):
        raise RuntimeError("compute job byte length changed")
    return job


def _read_timeline(path: Path, expected_count: int, start: float, end: float) -> list[float]:
    timestamps: list[float] = []
    with path.open("r", encoding="utf-8") as stream:
        for expected_index, line in enumerate(stream):
            value = json.loads(line)
            if not isinstance(value, dict) or value.get("frame_index") != expected_index:
                raise RuntimeError("decoded frame timeline index changed")
            session_seconds = value.get("session_seconds")
            if not isinstance(session_seconds, (int, float)):
                raise RuntimeError("decoded frame timestamp is invalid")
            timestamp = float(session_seconds)
            if timestamps and timestamp <= timestamps[-1]:
                raise RuntimeError("decoded frame timestamps are not increasing")
            timestamps.append(timestamp)
    if len(timestamps) != expected_count:
        raise RuntimeError("decoded frame count differs from the camera job")
    if timestamps[0] < start - 0.001 or timestamps[-1] > end + 0.001:
        raise RuntimeError("decoded frame timeline escapes the camera epoch")
    return timestamps


def _palette(index: int) -> tuple[int, int, int]:
    digest = hashlib.sha256(f"mission-core-segment-{index}".encode()).digest()
    return (64 + digest[0] % 176, 64 + digest[1] % 176, 64 + digest[2] % 176)


def _blend(image: Any, colors: Any, alpha: float) -> Any:
    import numpy as np

    return np.clip(
        image.astype(np.float32) * (1.0 - alpha) + colors.astype(np.float32) * alpha,
        0,
        255,
    ).astype(np.uint8)


def _semantic_overlay(
    image: Any,
    index_map: Any,
    labels: dict[int, str],
) -> tuple[Any, list[dict[str, Any]]]:
    import numpy as np

    colors = np.zeros_like(image)
    counts = np.bincount(index_map.reshape(-1), minlength=max(labels) + 1)
    present = np.flatnonzero(counts)
    for index in present:
        colors[index_map == index] = _palette(int(index))
    total = int(index_map.size)
    classes = [
        {
            "id": int(index),
            "label": labels.get(int(index), f"class-{int(index)}"),
            "pixels": int(counts[index]),
            "fraction": round(float(counts[index]) / total, 9),
        }
        for index in sorted(present, key=lambda value: int(counts[value]), reverse=True)
    ]
    return _blend(image, colors, SEMANTIC_ALPHA), classes


def _instance_overlay(image: Any, index_map: Any, instances: list[dict[str, Any]]) -> Any:
    import numpy as np
    from PIL import Image, ImageDraw

    result = image.copy()
    for item in instances:
        instance_id = int(item["instance_id"])
        mask = index_map == instance_id
        color = np.asarray(_palette(500 + instance_id), dtype=np.uint8)
        if bool(mask.any()):
            result[mask] = _blend(
                result[mask],
                np.broadcast_to(color, result[mask].shape),
                INSTANCE_ALPHA,
            )
    canvas = Image.fromarray(result)
    draw = ImageDraw.Draw(canvas)
    for item in instances:
        color = _palette(500 + int(item["instance_id"]))
        x1, y1, x2, y2 = (int(round(value)) for value in item["box_xyxy"])
        draw.rectangle((x1, y1, x2, y2), outline=color, width=2)
        label = f"{item['label']} {float(item['score']):.0%}"
        text_box = draw.textbbox((x1, max(0, y1 - 13)), label)
        draw.rectangle(text_box, fill=(7, 8, 10))
        draw.text((x1, max(0, y1 - 13)), label, fill=color)
    return np.asarray(canvas, dtype=np.uint8)


def _write_png(path: Path, array: Any) -> None:
    from PIL import Image

    Image.fromarray(array).save(path, format="PNG", optimize=False)


def _percentiles(values: list[float]) -> dict[str, float]:
    if not values:
        return {"mean": 0.0, "p50": 0.0, "p95": 0.0, "p99": 0.0, "max": 0.0}
    ordered = sorted(values)

    def percentile(fraction: float) -> float:
        index = min(len(ordered) - 1, max(0, round((len(ordered) - 1) * fraction)))
        return round(ordered[index], 6)

    return {
        "mean": round(statistics.fmean(values), 6),
        "p50": percentile(0.5),
        "p95": percentile(0.95),
        "p99": percentile(0.99),
        "max": round(max(values), 6),
    }


class _GpuTelemetry:
    def __init__(self, stream: TextIO, interval_seconds: float) -> None:
        if not 0.25 <= interval_seconds <= 60:
            raise RuntimeError("telemetry interval is outside bounds")
        self._stream = stream
        self._interval = interval_seconds
        self._stop = threading.Event()
        self._thread = threading.Thread(target=self._run, name="gpu-telemetry", daemon=True)
        self.samples: list[dict[str, float]] = []

    def __enter__(self) -> _GpuTelemetry:
        self._thread.start()
        return self

    def __exit__(self, *_: object) -> None:
        self._stop.set()
        self._thread.join(timeout=self._interval + 5)

    def _run(self) -> None:
        query = (
            "timestamp,utilization.gpu,utilization.memory,memory.used,memory.total,"
            "temperature.gpu,power.draw"
        )
        while not self._stop.is_set():
            before = time.time()
            try:
                completed = subprocess.run(
                    [
                        "nvidia-smi",
                        f"--query-gpu={query}",
                        "--format=csv,noheader,nounits",
                    ],
                    capture_output=True,
                    check=True,
                    text=True,
                    timeout=5,
                )
                values = [item.strip() for item in completed.stdout.splitlines()[0].split(",")]
                sample = {
                    "epoch_seconds": round(before, 6),
                    "gpu_utilization_percent": float(values[1]),
                    "gpu_memory_utilization_percent": float(values[2]),
                    "gpu_memory_used_mib": float(values[3]),
                    "gpu_memory_total_mib": float(values[4]),
                    "gpu_temperature_c": float(values[5]),
                    "gpu_power_w": float(values[6]),
                }
                self.samples.append(sample)
                self._stream.write(json.dumps(sample, sort_keys=True) + "\n")
                self._stream.flush()
            except (OSError, subprocess.SubprocessError, ValueError, IndexError):
                pass
            self._stop.wait(max(0.0, self._interval - (time.time() - before)))

    def summary(self) -> dict[str, object]:
        fields = (
            "gpu_utilization_percent",
            "gpu_memory_utilization_percent",
            "gpu_memory_used_mib",
            "gpu_temperature_c",
            "gpu_power_w",
        )
        return {
            "sample_count": len(self.samples),
            "interval_seconds": self._interval,
            **{
                field: _percentiles([float(sample[field]) for sample in self.samples])
                for field in fields
            },
        }


def _model_files(semantic_snapshot: Path, checkpoint_path: Path) -> list[dict[str, object]]:
    files = [
        {"name": f"beit/{path.name}", "bytes": path.stat().st_size, "sha256": _sha256(path)}
        for path in sorted(semantic_snapshot.iterdir())
        if path.is_file()
    ]
    files.append(
        {
            "name": f"torchvision/{checkpoint_path.name}",
            "bytes": checkpoint_path.stat().st_size,
            "sha256": _sha256(checkpoint_path),
        }
    )
    return files


def _preflight(args: argparse.Namespace) -> int:
    job = _validate_job(args.job.resolve(strict=True))
    cache_root = args.cache.resolve(strict=True)

    import torch
    import torchvision  # noqa: F401
    import transformers  # noqa: F401
    from huggingface_hub import snapshot_download
    from torchvision.models.detection import MaskRCNN_ResNet50_FPN_V2_Weights

    if not torch.cuda.is_available() or torch.cuda.device_count() < 1:
        raise RuntimeError("CUDA device 0 is unavailable")
    cuda_probe = torch.zeros(1, device="cuda:0")
    torch.cuda.synchronize()
    del cuda_probe
    checkpoint = Path(MaskRCNN_ResNet50_FPN_V2_Weights.DEFAULT.url).name
    checkpoint_path = Path(torch.hub.get_dir()) / "checkpoints" / checkpoint
    if not checkpoint_path.is_file() or checkpoint_path.stat().st_size < 1:
        raise RuntimeError("cached instance checkpoint is unavailable")
    semantic_snapshot = Path(
        snapshot_download(
            repo_id=SEMANTIC_MODEL_ID,
            revision=SEMANTIC_REVISION,
            cache_dir=cache_root / "huggingface",
            allow_patterns=("config.json", "preprocessor_config.json", "pytorch_model.bin"),
        )
    )
    required = ("config.json", "preprocessor_config.json", "pytorch_model.bin")
    if any(not (semantic_snapshot / name).is_file() for name in required):
        raise RuntimeError("cached semantic checkpoint is incomplete")
    print(
        json.dumps(
            {
                "state": "preflight-ready",
                "job_id": job["job_id"],
                "cuda_device": torch.cuda.get_device_name(),
            },
            sort_keys=True,
        ),
        flush=True,
    )
    return 0


def _run(args: argparse.Namespace) -> int:
    if not _valid_sha256(args.calibration_sha256):
        raise RuntimeError("calibration SHA-256 is invalid")
    job_root = args.job.resolve(strict=True)
    frames_root = args.frames.resolve(strict=True)
    timeline_path = args.timeline.resolve(strict=True)
    output_root = args.output.resolve()
    cache_root = args.cache.resolve()
    if output_root.exists() or not frames_root.is_dir():
        raise RuntimeError("output must be absent and frames must be a directory")
    job = _validate_job(job_root)
    input_document = job["input"]
    timeline = input_document["timeline"]
    segment_count = int(input_document["segment_count"])
    frame_paths = [frames_root / f"frame-{index:06d}.png" for index in range(1, segment_count + 1)]
    if not all(path.is_file() for path in frame_paths):
        raise RuntimeError("decoded frame set is incomplete")
    if len(list(frames_root.glob("frame-*.png"))) != segment_count:
        raise RuntimeError("decoded frame set contains unexpected files")
    timestamps = _read_timeline(
        timeline_path,
        segment_count,
        float(timeline["start_seconds"]),
        float(timeline["end_seconds"]),
    )

    import numpy as np
    import torch
    import torch.nn.functional as functional
    import torchvision
    import transformers
    from huggingface_hub import snapshot_download
    from PIL import Image
    from torchvision.models.detection import (
        MaskRCNN_ResNet50_FPN_V2_Weights,
        maskrcnn_resnet50_fpn_v2,
    )
    from transformers import AutoImageProcessor, BeitForSemanticSegmentation

    if not torch.cuda.is_available():
        raise RuntimeError("CUDA is required for a perception epoch run")
    cuda_index = 0
    device = torch.device(f"cuda:{cuda_index}")
    cache_root.mkdir(mode=0o700, parents=True, exist_ok=True)
    output_root.mkdir(mode=0o700, parents=True, exist_ok=False)
    instance_root = output_root / "instance-masks"
    semantic_root = output_root / "semantic-masks"
    overlay_root = output_root / "overlay-frames"
    instance_root.mkdir(mode=0o700)
    semantic_root.mkdir(mode=0o700)
    overlay_root.mkdir(mode=0o700)
    telemetry_path = output_root / "gpu-telemetry.jsonl"
    started = time.perf_counter()
    instance_latencies: list[float] = []
    semantic_latencies: list[float] = []
    instances_by_frame: list[list[dict[str, Any]]] = []
    total_instances = 0

    with telemetry_path.open("x", encoding="utf-8", newline="\n") as telemetry_stream:
        with _GpuTelemetry(telemetry_stream, args.telemetry_interval_seconds) as telemetry:
            instance_weights = MaskRCNN_ResNet50_FPN_V2_Weights.DEFAULT
            instance_model = maskrcnn_resnet50_fpn_v2(weights=instance_weights).to(device).eval()
            instance_labels = list(instance_weights.meta["categories"])
            with torch.inference_mode():
                for frame_index, path in enumerate(frame_paths):
                    with Image.open(path) as opened:
                        image = np.asarray(opened.convert("RGB"), dtype=np.uint8)
                    tensor = instance_weights.transforms()(Image.fromarray(image)).to(device)
                    torch.cuda.synchronize()
                    before = time.perf_counter()
                    prediction = instance_model([tensor])[0]
                    torch.cuda.synchronize()
                    instance_latencies.append((time.perf_counter() - before) * 1000.0)
                    scores = prediction["scores"].detach().cpu().numpy()
                    keep = np.flatnonzero(scores >= INSTANCE_SCORE_THRESHOLD)
                    index_map = np.zeros(image.shape[:2], dtype=np.uint16)
                    instances: list[dict[str, Any]] = []
                    for output_index in keep:
                        mask = (
                            prediction["masks"][output_index, 0].detach().cpu().numpy()
                            >= INSTANCE_MASK_THRESHOLD
                        )
                        mask_pixels = int(mask.sum())
                        if mask_pixels < 8:
                            continue
                        label_id = int(prediction["labels"][output_index].item())
                        instance_id = len(instances) + 1
                        index_map[np.logical_and(mask, index_map == 0)] = instance_id
                        instances.append(
                            {
                                "instance_id": instance_id,
                                "class_id": label_id,
                                "label": instance_labels[label_id],
                                "score": round(float(scores[output_index]), 9),
                                "box_xyxy": [
                                    round(float(value), 6)
                                    for value in prediction["boxes"][output_index]
                                    .detach()
                                    .cpu()
                                    .tolist()
                                ],
                                "mask_pixels": mask_pixels,
                            }
                        )
                    _write_png(instance_root / f"frame-{frame_index + 1:06d}.png", index_map)
                    instances_by_frame.append(instances)
                    total_instances += len(instances)
                    if (frame_index + 1) % 100 == 0 or frame_index + 1 == segment_count:
                        print(
                            json.dumps(
                                {
                                    "phase": "instance",
                                    "frames_processed": frame_index + 1,
                                    "frames_total": segment_count,
                                },
                                sort_keys=True,
                            ),
                            flush=True,
                        )
            checkpoint = Path(instance_weights.url).name
            checkpoint_path = Path(torch.hub.get_dir()) / "checkpoints" / checkpoint
            del instance_model
            torch.cuda.empty_cache()

            semantic_snapshot = Path(
                snapshot_download(
                    repo_id=SEMANTIC_MODEL_ID,
                    revision=SEMANTIC_REVISION,
                    cache_dir=cache_root / "huggingface",
                    allow_patterns=("config.json", "preprocessor_config.json", "pytorch_model.bin"),
                )
            )
            processor = AutoImageProcessor.from_pretrained(
                semantic_snapshot,
                local_files_only=True,
                use_fast=False,
            )
            semantic_model, loading = BeitForSemanticSegmentation.from_pretrained(
                semantic_snapshot,
                local_files_only=True,
                output_loading_info=True,
            )
            load_problems = {
                name: loading.get(name, [])
                for name in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs")
                if loading.get(name)
            }
            if load_problems:
                raise RuntimeError(
                    "semantic checkpoint did not load exactly: "
                    + json.dumps(load_problems)
                )
            semantic_model = semantic_model.to(device).eval()
            semantic_labels = {
                int(key): str(value)
                for key, value in semantic_model.config.id2label.items()
            }
            frame_metadata = output_root / "frames.jsonl"
            with frame_metadata.open("x", encoding="utf-8", newline="\n") as metadata_stream:
                with torch.inference_mode():
                    for frame_index, (path, session_seconds) in enumerate(
                        zip(frame_paths, timestamps, strict=True)
                    ):
                        with Image.open(path) as opened:
                            image = np.asarray(opened.convert("RGB"), dtype=np.uint8)
                        inputs = processor(images=Image.fromarray(image), return_tensors="pt")
                        inputs = {name: value.to(device) for name, value in inputs.items()}
                        torch.cuda.synchronize()
                        before = time.perf_counter()
                        logits = semantic_model(**inputs).logits
                        resized = functional.interpolate(
                            logits,
                            size=image.shape[:2],
                            mode="bilinear",
                            align_corners=False,
                        )
                        semantic = resized.argmax(dim=1)[0].detach().cpu().numpy().astype(np.uint8)
                        torch.cuda.synchronize()
                        semantic_latencies.append((time.perf_counter() - before) * 1000.0)
                        semantic_overlay, classes = _semantic_overlay(
                            image,
                            semantic,
                            semantic_labels,
                        )
                        instance_map = np.asarray(
                            Image.open(instance_root / f"frame-{frame_index + 1:06d}.png"),
                            dtype=np.uint16,
                        )
                        overlay = _instance_overlay(
                            semantic_overlay,
                            instance_map,
                            instances_by_frame[frame_index],
                        )
                        _write_png(semantic_root / f"frame-{frame_index + 1:06d}.png", semantic)
                        _write_png(overlay_root / f"frame-{frame_index + 1:06d}.png", overlay)
                        metadata_stream.write(
                            json.dumps(
                                {
                                    "schema_version": FRAME_SCHEMA,
                                    "frame_index": frame_index,
                                    "sequence": frame_index + 1,
                                    "session_seconds": round(session_seconds, 9),
                                    "instances": instances_by_frame[frame_index],
                                    "semantic_classes": classes,
                                },
                                ensure_ascii=False,
                                sort_keys=True,
                                separators=(",", ":"),
                                allow_nan=False,
                            )
                            + "\n"
                        )
                        if (frame_index + 1) % 100 == 0 or frame_index + 1 == segment_count:
                            print(
                                json.dumps(
                                    {
                                        "phase": "semantic",
                                        "frames_processed": frame_index + 1,
                                        "frames_total": segment_count,
                                    },
                                    sort_keys=True,
                                ),
                                flush=True,
                            )
                metadata_stream.flush()
                os.fsync(metadata_stream.fileno())
            del semantic_model
            torch.cuda.empty_cache()

        inference_elapsed = time.perf_counter() - started
        model_files = _model_files(semantic_snapshot, checkpoint_path)
        report = {
            "schema_version": REPORT_SCHEMA,
            "run_id": f"{job['job_id']}-panoptic-v1",
            "created_at_utc": datetime.now(UTC)
            .isoformat(timespec="milliseconds")
            .replace("+00:00", "Z"),
            "state": "inference-complete-awaiting-publication",
            "input": {
                "job_id": job["job_id"],
                "input_sha256": job["input_sha256"],
                "session_id": input_document["session_id"],
                "source_id": input_document["source_id"],
                "codec_epoch": input_document["codec_epoch"],
                "segment_count": segment_count,
                "byte_length": input_document["byte_length"],
                "timeline_start_seconds": timeline["start_seconds"],
                "timeline_end_seconds": timeline["end_seconds"],
                "frame_timeline_sha256": _sha256(timeline_path),
            },
            "calibration": {
                "content_identity_sha256": args.calibration_sha256,
                "camera_slot": args.calibration_slot,
            },
            "configuration": {
                "pipeline": "recorded-panoptic-maskrcnn-beit/v1",
                "instance_score_threshold": INSTANCE_SCORE_THRESHOLD,
                "instance_mask_threshold": INSTANCE_MASK_THRESHOLD,
                "semantic_alpha": SEMANTIC_ALPHA,
                "instance_alpha": INSTANCE_ALPHA,
                "batch_size": 1,
                "frame_policy": "all-frames-no-sampling",
                "runner_sha256": _sha256(Path(__file__).resolve(strict=True)),
            },
            "models": {
                "instance": {
                    "id": "torchvision/maskrcnn_resnet50_fpn_v2",
                    "weights": str(instance_weights),
                    "score_threshold": INSTANCE_SCORE_THRESHOLD,
                },
                "semantic": {
                    "id": SEMANTIC_MODEL_ID,
                    "revision": SEMANTIC_REVISION,
                    "checkpoint_load": "exact-no-missing-unexpected-or-mismatched-keys",
                },
                "files": model_files,
            },
            "runtime": {
                "hostname": platform.node(),
                "platform": platform.platform(),
                "python": platform.python_version(),
                "torch": torch.__version__,
                "torchvision": torchvision.__version__,
                "transformers": transformers.__version__,
                "cuda_runtime": torch.version.cuda,
                "gpu": torch.cuda.get_device_name(),
                "gpu_compute_capability": list(torch.cuda.get_device_capability()),
            },
            "metrics": {
                "frames_expected": segment_count,
                "frames_processed": segment_count,
                "frames_failed": 0,
                "frames_skipped": 0,
                "instances": total_instances,
                "inference_wall_seconds": round(inference_elapsed, 6),
                "inference_frames_per_second": round(segment_count / inference_elapsed, 6),
                "instance_latency_ms": _percentiles(instance_latencies),
                "semantic_latency_ms": _percentiles(semantic_latencies),
                "cuda_peak_memory_allocated_mib": round(
                    torch.cuda.max_memory_allocated() / 2**20,
                    3,
                ),
                "cuda_peak_memory_reserved_mib": round(
                    torch.cuda.max_memory_reserved() / 2**20,
                    3,
                ),
                "process_peak_rss_mib": round(
                    resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / 1024,
                    3,
                ),
                "system_load_average": [round(value, 6) for value in os.getloadavg()],
                "gpu_telemetry": telemetry.summary(),
            },
            "versions": {
                name: importlib.metadata.version(name)
                for name in ("numpy", "pillow", "torch", "torchvision", "transformers")
            },
        }
        _write_json(output_root / "run-report.partial.json", report)
    print(json.dumps({"state": "inference-complete", "frames": segment_count}, sort_keys=True))
    return 0


def _artifact(
    path: Path,
    kind: str,
    media_type: str,
    schema_version: str | None = None,
) -> dict[str, object]:
    result: dict[str, object] = {
        "kind": kind,
        "path": path.name,
        "media_type": media_type,
        "byte_length": path.stat().st_size,
        "sha256": _sha256(path),
    }
    if schema_version is not None:
        result["schema_version"] = schema_version
    return result


def _finalize(args: argparse.Namespace) -> int:
    output_root = args.output.resolve(strict=True)
    video = args.video.resolve(strict=True)
    masks = args.masks.resolve(strict=True)
    if video.parent != output_root or masks.parent != output_root:
        raise RuntimeError("published artifacts must be direct result children")
    partial_path = output_root / "run-report.partial.json"
    frames_path = output_root / "frames.jsonl"
    telemetry_path = output_root / "gpu-telemetry.jsonl"
    partial = _read_object(partial_path)
    if partial.get("schema_version") != REPORT_SCHEMA:
        raise RuntimeError("partial run report is incompatible")
    input_document = partial["input"]
    configuration = partial["configuration"]
    models = partial["models"]
    identity = {
        "schema_version": "missioncore.recorded-perception-identity/v2",
        "job_id": input_document["job_id"],
        "input_sha256": input_document["input_sha256"],
        "calibration": partial["calibration"],
        "configuration": configuration,
        "models": models,
        "publication": {
            "video_encoder": args.encoder,
            "video_media_type": "video/mp4",
            "mask_archive_media_type": "application/gzip",
        },
    }
    identity_sha256 = hashlib.sha256(_canonical_json(identity)).hexdigest()
    result_id = f"result-{identity_sha256}"
    final_report = {
        **partial,
        "state": "published",
        "result_id": result_id,
        "publication": {
            "extract_seconds": round(args.extract_seconds, 6),
            "encode_seconds": round(args.encode_seconds, 6),
            "wall_seconds": round(args.wall_seconds, 6),
            "end_to_end_frames_per_second": round(
                int(partial["metrics"]["frames_processed"]) / args.wall_seconds,
                6,
            ),
            "encoder": args.encoder,
        },
    }
    report_path = output_root / "run-report.json"
    _write_json(report_path, final_report)
    artifacts = [
        _artifact(video, "panoptic-overlay-video", "video/mp4"),
        _artifact(masks, "panoptic-mask-archive", "application/gzip"),
        _artifact(frames_path, "panoptic-frame-metadata", "application/x-ndjson", FRAME_SCHEMA),
        _artifact(telemetry_path, "worker-gpu-telemetry", "application/x-ndjson"),
        _artifact(report_path, "perception-run-report", "application/json", REPORT_SCHEMA),
    ]
    result = {
        "schema_version": RESULT_SCHEMA,
        "result_id": result_id,
        "identity_sha256": identity_sha256,
        "identity": identity,
        "created_at_utc": final_report["created_at_utc"],
        "job_id": input_document["job_id"],
        "input_sha256": input_document["input_sha256"],
        "session_id": input_document["session_id"],
        "source_id": input_document["source_id"],
        "codec_epoch": input_document["codec_epoch"],
        "timestamp_basis": "session-time-seconds",
        "timeline_start_seconds": input_document["timeline_start_seconds"],
        "timeline_end_seconds": input_document["timeline_end_seconds"],
        "frames_processed": final_report["metrics"]["frames_processed"],
        "artifacts": artifacts,
    }
    _write_json(output_root / "result.json", result)
    partial_path.unlink()
    print(json.dumps({"result_id": result_id, "identity_sha256": identity_sha256}, sort_keys=True))
    return 0


def main() -> int:
    args = _arguments()
    if args.command == "preflight":
        return _preflight(args)
    if args.command == "run":
        return _run(args)
    return _finalize(args)


if __name__ == "__main__":
    raise SystemExit(main())
