#!/usr/bin/env python3
"""Build an immutable CVAT review handoff for the LAB E3 control and challenger."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
import zipfile
from pathlib import Path, PurePosixPath
from typing import Any

import numpy as np
from PIL import Image

RESULT_SCHEMA = "missioncore.k1-e3-rectified-segmentation-result/v1"
PACK_SCHEMA = "missioncore.perception-evaluation-pack/v1"
LABEL_SCHEMA = "missioncore.perception-cvat-labels/v1"
MANIFEST_SCHEMA = "missioncore.lab-e3-cvat-review-workspace/v1"
IDENTITY_SCHEMA = "missioncore.lab-e3-cvat-review-identity/v1"
CONTROL = "eomt-fisheye-mask"
CHALLENGER = "eomt-kb4-cubemap5-clahe"
EXPECTED_FRAMES = 64
PRIORITY_IMAGE_IDS = (59, 63, 62, 61, 60, 64, 50, 30, 17, 29)
ZIP_TIMESTAMP = (1980, 1, 1, 0, 0, 0)


def _arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--evaluation-pack", type=Path, required=True)
    parser.add_argument("--e3-result", type=Path, required=True)
    parser.add_argument("--labels", type=Path, required=True)
    parser.add_argument("--images-zip", type=Path, required=True)
    parser.add_argument("--output-root", type=Path, required=True)
    return parser.parse_args()


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


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


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


def _valid_sha256(value: object) -> bool:
    return (
        isinstance(value, str)
        and len(value) == 64
        and all(character in "0123456789abcdef" for character in value)
    )


def _safe_artifact(root: Path, encoded: object) -> Path:
    if not isinstance(encoded, str):
        raise RuntimeError("artifact path is not a string")
    relative = PurePosixPath(encoded)
    if relative.is_absolute() or ".." in relative.parts or not relative.parts:
        raise RuntimeError("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("artifact path escaped its root")
    return path


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


def _validate_pack(root: Path) -> tuple[dict[str, Any], list[dict[str, Any]]]:
    manifest = _read_object(root / "manifest.json")
    identity = manifest.get("identity")
    identity_sha256 = manifest.get("identity_sha256")
    if (
        manifest.get("schema_version") != PACK_SCHEMA
        or not isinstance(identity, dict)
        or not _valid_sha256(identity_sha256)
        or hashlib.sha256(_canonical_json(identity)).hexdigest() != identity_sha256
        or manifest.get("generation_id") != f"evaluation-pack-{identity_sha256}"
        or identity.get("source_id") != "sensor.camera.right"
        or identity.get("calibration_slot") != "camera_1"
    ):
        raise RuntimeError("evaluation pack is incompatible")
    frames = identity.get("frames")
    if (
        not isinstance(frames, list)
        or len(frames) != EXPECTED_FRAMES
        or [frame.get("image_id") for frame in frames] != list(range(1, 65))
    ):
        raise RuntimeError("evaluation pack frame contract changed")
    return manifest, frames


def _validate_result(root: Path, pack: dict[str, Any]) -> dict[str, Any]:
    result = _read_object(root / "result.json")
    identity = result.get("identity")
    identity_sha256 = result.get("identity_sha256")
    if (
        result.get("schema_version") != RESULT_SCHEMA
        or result.get("ground_truth") is not False
        or not isinstance(identity, dict)
        or not _valid_sha256(identity_sha256)
        or hashlib.sha256(_canonical_json(identity)).hexdigest() != identity_sha256
        or result.get("result_id") != f"e3-segmentation-{identity_sha256}"
        or root.name != result.get("result_id")
        or identity.get("evaluation_pack_id") != pack["generation_id"]
        or identity.get("evaluation_identity_sha256") != pack["identity_sha256"]
        or identity.get("frame_count") != EXPECTED_FRAMES
        or CONTROL not in identity.get("variants", ())
        or CHALLENGER not in identity.get("variants", ())
    ):
        raise RuntimeError("LAB E3 result is incompatible")
    artifacts = result.get("artifacts")
    if not isinstance(artifacts, list):
        raise RuntimeError("LAB E3 artifact list is invalid")
    descriptors = {
        item.get("path"): item
        for item in artifacts
        if isinstance(item, dict) and isinstance(item.get("path"), str)
    }
    for variant in (CONTROL, CHALLENGER):
        for image_id in range(1, 65):
            encoded = f"semantic-masks/{variant}/image-{image_id:03d}.png"
            descriptor = descriptors.get(encoded)
            path = _safe_artifact(root, encoded)
            if (
                not isinstance(descriptor, dict)
                or path.stat().st_size != descriptor.get("bytes")
                or not _valid_sha256(descriptor.get("sha256"))
                or _sha256(path) != descriptor["sha256"]
            ):
                raise RuntimeError(f"LAB E3 semantic artifact changed: {encoded}")
    return result


def _semantic_labels(document: dict[str, Any]) -> list[dict[str, Any]]:
    if document.get("schema_version") != LABEL_SCHEMA:
        raise RuntimeError("CVAT label document is incompatible")
    raw = document.get("semantic_task")
    if not isinstance(raw, list):
        raise RuntimeError("CVAT semantic labels are absent")
    labels = [
        {"id": 0, "name": "background", "color": "#000000"},
        *[
            {"id": int(item["id"]), "name": str(item["name"]), "color": str(item["color"])}
            for item in raw
        ],
    ]
    if [label["id"] for label in labels] != list(range(16)) + [255]:
        raise RuntimeError("CVAT semantic taxonomy changed")
    return labels


def _rgb(value: str) -> tuple[int, int, int]:
    if len(value) != 7 or not value.startswith("#"):
        raise RuntimeError("CVAT label color is invalid")
    return tuple(int(value[index : index + 2], 16) for index in (1, 3, 5))


def _zip_info(name: str) -> zipfile.ZipInfo:
    info = zipfile.ZipInfo(name, date_time=ZIP_TIMESTAMP)
    info.compress_type = zipfile.ZIP_DEFLATED
    info.external_attr = 0o100644 << 16
    return info


def _write_segmentation_zip(
    destination: Path,
    *,
    result_root: Path,
    variant: str,
    frames: list[dict[str, Any]],
    labels: list[dict[str, Any]],
) -> None:
    palette = np.zeros((256, 3), dtype=np.uint8)
    for label in labels:
        palette[label["id"]] = _rgb(label["color"])
    labelmap = "".join(
        f"{label['name']}:{','.join(str(value) for value in _rgb(label['color']))}::\n"
        for label in labels
    ).encode()
    stems = [
        f"image-{frame['image_id']:03d}-frame-{frame['frame_index']:06d}"
        for frame in frames
    ]
    with zipfile.ZipFile(destination, mode="x", compression=zipfile.ZIP_DEFLATED) as archive:
        archive.writestr(_zip_info("labelmap.txt"), labelmap)
        archive.writestr(
            _zip_info("ImageSets/Segmentation/default.txt"),
            ("\n".join(stems) + "\n").encode(),
        )
        for image_id, stem in enumerate(stems, start=1):
            source = result_root / "semantic-masks" / variant / f"image-{image_id:03d}.png"
            with Image.open(source) as opened:
                indexed = np.asarray(opened.convert("L"), dtype=np.uint8)
            if indexed.shape != (600, 800) or np.any(~np.isin(indexed, list(range(16)))):
                raise RuntimeError(f"semantic mask taxonomy changed: {source}")
            rgb = palette[indexed]
            from io import BytesIO

            stream = BytesIO()
            Image.fromarray(rgb, mode="RGB").save(stream, format="PNG", optimize=True)
            archive.writestr(
                _zip_info(f"SegmentationClass/{stem}.png"),
                stream.getvalue(),
            )


def _validate_existing(root: Path, identity_sha256: str) -> dict[str, Any]:
    manifest = _read_object(root / "manifest.json")
    if (
        manifest.get("schema_version") != MANIFEST_SCHEMA
        or manifest.get("workspace_id") != f"e3-cvat-review-{identity_sha256}"
        or manifest.get("identity_sha256") != identity_sha256
        or manifest.get("ground_truth") is not False
    ):
        raise RuntimeError("existing LAB E3 CVAT workspace is incompatible")
    for descriptor in manifest.get("artifacts", ()):
        path = _safe_artifact(root, descriptor.get("path"))
        if path.stat().st_size != descriptor.get("bytes") or _sha256(path) != descriptor.get(
            "sha256"
        ):
            raise RuntimeError("existing LAB E3 CVAT artifact changed")
    return manifest


def main() -> int:
    args = _arguments()
    pack_root = args.evaluation_pack.resolve(strict=True)
    result_root = args.e3_result.resolve(strict=True)
    labels_path = args.labels.resolve(strict=True)
    images_zip = args.images_zip.resolve(strict=True)
    output_root = args.output_root.resolve()
    pack, frames = _validate_pack(pack_root)
    result = _validate_result(result_root, pack)
    labels_document = _read_object(labels_path)
    labels = _semantic_labels(labels_document)
    with zipfile.ZipFile(images_zip) as archive:
        if len(archive.namelist()) != EXPECTED_FRAMES:
            raise RuntimeError("CVAT image archive frame count changed")

    identity = {
        "schema_version": IDENTITY_SCHEMA,
        "evaluation_pack_id": pack["generation_id"],
        "evaluation_identity_sha256": pack["identity_sha256"],
        "e3_result_id": result["result_id"],
        "e3_identity_sha256": result["identity_sha256"],
        "control_variant": CONTROL,
        "challenger_variant": CHALLENGER,
        "frame_count": EXPECTED_FRAMES,
        "labels_sha256": _sha256(labels_path),
        "images_zip_sha256": _sha256(images_zip),
        "priority_image_ids": list(PRIORITY_IMAGE_IDS),
        "producer_sha256": _sha256(Path(__file__).resolve(strict=True)),
    }
    identity_sha256 = hashlib.sha256(_canonical_json(identity)).hexdigest()
    workspace_id = f"e3-cvat-review-{identity_sha256}"
    final_root = output_root / workspace_id
    if final_root.exists():
        manifest = _validate_existing(final_root.resolve(strict=True), identity_sha256)
        print(json.dumps({"state": "reused", **manifest}, sort_keys=True))
        return 0

    output_root.mkdir(parents=True, exist_ok=True)
    staging = output_root / f".{workspace_id}.{os.getpid()}.publish"
    staging.mkdir(mode=0o700)
    try:
        labels_output = staging / "labels.json"
        labels_output.write_text(json.dumps(labels_document, indent=2) + "\n", encoding="utf-8")
        for variant, filename in (
            (CONTROL, "control-fisheye-mask.zip"),
            (CHALLENGER, "challenger-kb4-cubemap5.zip"),
        ):
            _write_segmentation_zip(
                staging / filename,
                result_root=result_root,
                variant=variant,
                frames=frames,
                labels=labels,
            )
        artifacts = [
            _artifact(path, staging)
            for path in sorted(staging.iterdir())
            if path.is_file()
        ]
        manifest = {
            "schema_version": MANIFEST_SCHEMA,
            "workspace_id": workspace_id,
            "identity_sha256": identity_sha256,
            "identity": identity,
            "ground_truth": False,
            "artifacts": artifacts,
            "next_gate": "two-pass CVAT review and reviewed export",
        }
        (staging / "manifest.json").write_text(
            json.dumps(manifest, indent=2) + "\n",
            encoding="utf-8",
        )
        staging.replace(final_root)
    finally:
        if staging.exists():
            shutil.rmtree(staging)
    print(
        json.dumps(
            {"state": "created", "workspace_id": workspace_id, "root": str(final_root)},
            sort_keys=True,
        )
    )
    return 0


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