#!/usr/bin/env python3
"""Run a bounded, recorded-only segmentation probe on an external GPU worker."""

from __future__ import annotations

import argparse
import hashlib
import importlib.metadata
import json
import os
import platform
import shutil
import statistics
import time
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

MAX_INPUT_IMAGES = 16
MAX_INPUT_BYTES = 32 * 1024 * 1024
MAX_DIMENSION = 4096
INSTANCE_SCORE_THRESHOLD = 0.5
SEMANTIC_MODEL_ID = "microsoft/beit-base-finetuned-ade-640-640"
SEMANTIC_REVISION = "a8b6f5ef4acb2ea55d882989deaa02d39401e2b2"
RESULT_SCHEMA = "missioncore.recorded-segmentation-experiment/v1"


def _arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--input", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--cache", type=Path, 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 _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, 0.46), classes


def _instance_overlay(image: Any, instances: list[dict[str, Any]], masks: list[Any]) -> Any:
    import numpy as np

    result = image.copy()
    for index, (instance, mask) in enumerate(zip(instances, masks, strict=True), start=1):
        color = np.asarray(_palette(500 + index), dtype=np.uint8)
        result[mask] = _blend(result[mask], np.broadcast_to(color, result[mask].shape), 0.56)
        x1, y1, x2, y2 = (int(round(value)) for value in instance["box_xyxy"])
        result[max(0, y1) : min(result.shape[0], y1 + 2), max(0, x1) : min(result.shape[1], x2)] = (
            color
        )
        result[max(0, y2 - 2) : min(result.shape[0], y2), max(0, x1) : min(result.shape[1], x2)] = (
            color
        )
        result[max(0, y1) : min(result.shape[0], y2), max(0, x1) : min(result.shape[1], x1 + 2)] = (
            color
        )
        result[max(0, y1) : min(result.shape[0], y2), max(0, x2 - 2) : min(result.shape[1], x2)] = (
            color
        )
    return result


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

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


def _artifact(path: Path) -> dict[str, object]:
    return {"name": path.name, "bytes": path.stat().st_size, "sha256": _sha256(path)}


def main() -> int:
    args = _arguments()
    input_root = args.input.expanduser().resolve(strict=True)
    output_root = args.output.expanduser().resolve()
    cache_root = args.cache.expanduser().resolve()
    if not input_root.is_dir() or output_root.exists():
        raise RuntimeError("input must be a directory and output must not exist")
    images = sorted(input_root.glob("camera-*.png"))
    if not images or len(images) > MAX_INPUT_IMAGES:
        raise RuntimeError("input image count is outside the recorded probe bound")
    for path in images:
        if not path.is_file() or path.stat().st_size > MAX_INPUT_BYTES:
            raise RuntimeError("input image is missing or too large")

    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 this worker experiment")
    device = torch.device("cuda:0")
    cache_root.mkdir(mode=0o700, parents=True, exist_ok=True)
    staging = output_root.with_name(f".{output_root.name}.incomplete")
    staging.mkdir(mode=0o700, parents=True, exist_ok=False)
    started = time.perf_counter()
    try:
        decoded: list[tuple[Path, Any]] = []
        width: int | None = None
        height: int | None = None
        for path in images:
            with Image.open(path) as opened:
                opened.verify()
            with Image.open(path) as opened:
                rgb = opened.convert("RGB")
                if max(rgb.size) > MAX_DIMENSION:
                    raise RuntimeError("input image dimension is outside the probe bound")
                if width is None:
                    width, height = rgb.size
                elif rgb.size != (width, height):
                    raise RuntimeError("input images do not share one physical epoch resolution")
                decoded.append((path, np.asarray(rgb, dtype=np.uint8)))

        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"])
        instance_results: dict[str, dict[str, Any]] = {}
        instance_latencies: list[float] = []
        with torch.inference_mode():
            for path, image in decoded:
                tensor = instance_weights.transforms()(Image.fromarray(image)).to(device)
                torch.cuda.synchronize()
                before = time.perf_counter()
                prediction = instance_model([tensor])[0]
                torch.cuda.synchronize()
                latency_ms = (time.perf_counter() - before) * 1000.0
                instance_latencies.append(latency_ms)
                scores = prediction["scores"].detach().cpu().numpy()
                keep = np.flatnonzero(scores >= INSTANCE_SCORE_THRESHOLD)
                masks: list[Any] = []
                instances: list[dict[str, Any]] = []
                index_map = np.zeros(image.shape[:2], dtype=np.uint16)
                for output_index in keep:
                    mask = prediction["masks"][output_index, 0].detach().cpu().numpy() >= 0.5
                    if int(mask.sum()) < 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
                    masks.append(mask)
                    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": int(mask.sum()),
                        }
                    )
                stem = path.stem
                _write_png(staging / f"{stem}.instances.png", index_map)
                _write_png(
                    staging / f"{stem}.instance-overlay.png",
                    _instance_overlay(image, instances, masks),
                )
                instance_results[stem] = {
                    "latency_ms": round(latency_ms, 6),
                    "instances": instances,
                }
        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, semantic_loading = BeitForSemanticSegmentation.from_pretrained(
            semantic_snapshot,
            local_files_only=True,
            output_loading_info=True,
        )
        semantic_load_problems = {
            name: semantic_loading.get(name, [])
            for name in ("missing_keys", "unexpected_keys", "mismatched_keys", "error_msgs")
            if semantic_loading.get(name)
        }
        if semantic_load_problems:
            raise RuntimeError(
                "semantic checkpoint did not load exactly: "
                + json.dumps(semantic_load_problems, sort_keys=True)
            )
        semantic_model = semantic_model.to(device).eval()
        semantic_labels = {
            int(key): str(value) for key, value in semantic_model.config.id2label.items()
        }
        semantic_results: dict[str, dict[str, Any]] = {}
        semantic_latencies: list[float] = []
        with torch.inference_mode():
            for path, image in decoded:
                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()
                latency_ms = (time.perf_counter() - before) * 1000.0
                semantic_latencies.append(latency_ms)
                overlay, classes = _semantic_overlay(image, semantic, semantic_labels)
                stem = path.stem
                _write_png(staging / f"{stem}.semantic.png", semantic)
                _write_png(staging / f"{stem}.semantic-overlay.png", overlay)
                semantic_results[stem] = {
                    "latency_ms": round(latency_ms, 6),
                    "classes": classes,
                }
        del semantic_model
        torch.cuda.empty_cache()

        rows = []
        for path, image in decoded:
            stem = path.stem
            instance_overlay = np.asarray(Image.open(staging / f"{stem}.instance-overlay.png"))
            semantic_overlay = np.asarray(Image.open(staging / f"{stem}.semantic-overlay.png"))
            rows.append(np.concatenate((image, instance_overlay, semantic_overlay), axis=1))
        _write_png(staging / "segmentation-mosaic.png", np.concatenate(rows, axis=0))

        checkpoint = Path(instance_weights.url).name
        checkpoint_path = Path(torch.hub.get_dir()) / "checkpoints" / checkpoint
        model_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()
        ]
        model_files.append(
            {
                "name": f"torchvision/{checkpoint}",
                "bytes": checkpoint_path.stat().st_size,
                "sha256": _sha256(checkpoint_path),
            }
        )
        inputs_manifest = [
            {"name": path.name, "bytes": path.stat().st_size, "sha256": _sha256(path)}
            for path in images
        ]
        identity = {
            "inputs": inputs_manifest,
            "instance_model": "torchvision/maskrcnn_resnet50_fpn_v2/default",
            "instance_threshold": INSTANCE_SCORE_THRESHOLD,
            "semantic_model": SEMANTIC_MODEL_ID,
            "semantic_revision": SEMANTIC_REVISION,
            "model_files": model_files,
        }
        manifest: dict[str, Any] = {
            "schema_version": RESULT_SCHEMA,
            "created_at_utc": datetime.now(UTC)
            .isoformat(timespec="milliseconds")
            .replace("+00:00", "Z"),
            "classification": "private-derived-recorded-perception-experiment",
            "generation_sha256": hashlib.sha256(_canonical_json(identity)).hexdigest(),
            "input": inputs_manifest,
            "models": {
                "instance": {
                    "id": "torchvision/maskrcnn_resnet50_fpn_v2",
                    "weights": str(instance_weights),
                    "weights_url": instance_weights.url,
                    "score_threshold": INSTANCE_SCORE_THRESHOLD,
                    "license": "BSD-3-Clause (TorchVision code/weight distribution)",
                },
                "semantic": {
                    "id": SEMANTIC_MODEL_ID,
                    "revision": SEMANTIC_REVISION,
                    "dataset": "ADE20K scene_parse_150",
                    "license": "Apache-2.0 (model-card metadata)",
                    "checkpoint_load": "exact-no-missing-unexpected-or-mismatched-keys",
                },
                "files": model_files,
            },
            "runtime": {
                "base_image": os.environ.get("MISSION_CORE_BASE_IMAGE", "unknown"),
                "python": platform.python_version(),
                "torch": torch.__version__,
                "torchvision": torchvision.__version__,
                "transformers": transformers.__version__,
                "cuda_runtime": torch.version.cuda,
                "cudnn": torch.backends.cudnn.version(),
                "gpu": torch.cuda.get_device_name(0),
                "packages": sorted(
                    (
                        {
                            "name": distribution.metadata["Name"],
                            "version": distribution.version,
                        }
                        for distribution in importlib.metadata.distributions()
                        if distribution.metadata["Name"]
                    ),
                    key=lambda item: str(item["name"]).lower(),
                ),
            },
            "metrics": {
                "frame_count": len(decoded),
                "instance_latency_ms": {
                    "mean": round(statistics.fmean(instance_latencies), 6),
                    "max": round(max(instance_latencies), 6),
                },
                "semantic_latency_ms": {
                    "mean": round(statistics.fmean(semantic_latencies), 6),
                    "max": round(max(semantic_latencies), 6),
                },
                "elapsed_seconds": round(time.perf_counter() - started, 6),
            },
            "frames": {
                stem: {
                    "instance": instance_results[stem],
                    "semantic": semantic_results[stem],
                }
                for stem in sorted(instance_results)
            },
            "acceptance": {
                "recorded_segmentation_artifact": "generated",
                "quality": "operator-review-required",
                "live": "not-tested",
                "safety": "not-accepted",
            },
        }
        artifacts = sorted(staging.glob("*.png"))
        manifest["outputs"] = [_artifact(path) for path in artifacts]
        (staging / "manifest.redacted.json").write_bytes(
            json.dumps(manifest, ensure_ascii=False, indent=2, sort_keys=True).encode("utf-8")
            + b"\n"
        )
        os.rename(staging, output_root)
    except BaseException:
        shutil.rmtree(staging, ignore_errors=True)
        raise
    print(
        json.dumps({"output": output_root.name, "generation_sha256": manifest["generation_sha256"]})
    )
    return 0


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