"""Operator review pack for a completed GOOSE ground qualification.

The public dataset and the full-resolution qualification cache remain on the
Simulation Worker D drive.  This derivative stores a bounded, point-aligned
preview for every validation frame so Mission Core can provide an honest
playback/review surface instead of exposing only aggregate metrics.
"""

from __future__ import annotations

import hashlib
import json
import math
import os
import tempfile
from concurrent.futures import ProcessPoolExecutor, as_completed
from pathlib import Path
from typing import Any, Final

import numpy as np

from k1link.datasets.goose_admission import (
    GOOSE_ARCHIVE_FILENAME,
    GOOSE_SOURCE_ID,
    GooseAdmissionError,
    read_goose_label_mapping,
)
from k1link.datasets.goose_profile import DEFAULT_GOOSE_PATCHWORK_PROFILE
from k1link.datasets.goose_qualification import (
    REPORT_ARTIFACT_KIND,
    _challenge_categories,
    _is_worker_dataset_root,
    _read_archive_frame,
    _sha256_file,
    _validation_frame_members,
)
from k1link.ground_segmentation import (
    DEFAULT_GROUND_BENCHMARK_PROFILE,
    GroundSegmenter,
    LocalPercentileGroundSegmenter,
    PatchworkPPGroundSegmenter,
)
from k1link.simulation import QualificationRunStore, RunState

GOOSE_REVIEW_PACK_SCHEMA: Final = "missioncore.goose-ground-review-pack/v1"
GOOSE_REVIEW_FRAME_SCHEMA: Final = "missioncore.goose-ground-review-frame/v1"
DEFAULT_REVIEW_POINTS: Final = 12_000
MAX_REVIEW_POINTS: Final = 20_000
_WORKER_PATCHWORK: GroundSegmenter | None = None


def build_goose_ground_review_pack(
    dataset_root: Path,
    runs_root: Path,
    *,
    run_id: str,
    parallel_workers: int = 8,
    preview_points: int = DEFAULT_REVIEW_POINTS,
    patchwork_module_name: str = "pypatchworkpp",
) -> dict[str, Any]:
    """Build or resume a compact visual review derivative for all run frames."""

    dataset = dataset_root.expanduser().absolute()
    runs = runs_root.expanduser().absolute()
    if not _is_worker_dataset_root(dataset):
        raise GooseAdmissionError("GOOSE review requires the canonical worker D root")
    if not 1 <= parallel_workers <= 32:
        raise GooseAdmissionError("parallel worker count is outside the admitted range")
    if not 1 <= preview_points <= MAX_REVIEW_POINTS:
        raise GooseAdmissionError("review point count is outside the admitted range")

    store = QualificationRunStore(runs, read_only=True)
    try:
        run = store.load(run_id)
    except Exception as exc:
        raise GooseAdmissionError("qualification run is unavailable") from exc
    if run.state is not RunState.COMPLETED:
        raise GooseAdmissionError("qualification run must be completed before review export")
    report_artifacts = [
        artifact for artifact in run.artifacts if artifact.kind == REPORT_ARTIFACT_KIND
    ]
    if len(report_artifacts) != 1:
        raise GooseAdmissionError("qualification run has no unique report artifact")
    report_artifact = report_artifacts[0]
    report_path = runs / run_id / report_artifact.relative_path
    if _sha256_file(report_path) != report_artifact.sha256:
        raise GooseAdmissionError("qualification report digest differs")
    try:
        report = json.loads(report_path.read_text(encoding="utf-8"))
    except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise GooseAdmissionError("qualification report is invalid") from exc
    if (
        not isinstance(report, dict)
        or report.get("run_id") != run_id
        or not isinstance(report.get("frames"), list)
    ):
        raise GooseAdmissionError("qualification report has no frame evidence")

    admission_path = dataset / "state/goose-3d-v2025-08-22.json"
    try:
        admission = json.loads(admission_path.read_text(encoding="utf-8"))
    except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise GooseAdmissionError("GOOSE admission manifest is unavailable") from exc
    archive_sha256 = str(admission.get("archive", {}).get("sha256", ""))
    if len(archive_sha256) != 64:
        raise GooseAdmissionError("GOOSE archive identity is invalid")
    install_root = dataset / "goose-3d/v2025-08-22/installs" / archive_sha256
    archive_path = dataset / "goose-3d/v2025-08-22/archives" / GOOSE_ARCHIVE_FILENAME
    label_mapping = read_goose_label_mapping(install_root / "goose_label_mapping.csv")
    members = _validation_frame_members(archive_path)
    report_frames = {
        str(frame["frame_id"]): frame
        for frame in report["frames"]
        if isinstance(frame, dict) and isinstance(frame.get("frame_id"), str)
    }
    if len(members) != len(report_frames):
        raise GooseAdmissionError("review source differs from the qualification frame set")

    identity_document = {
        "schema_version": GOOSE_REVIEW_PACK_SCHEMA,
        "source_run_id": run_id,
        "qualification_report_sha256": report_artifact.sha256,
        "archive_sha256": archive_sha256,
        "preview_points": preview_points,
        "sampling": "deterministic-even-index",
        "coordinate_codec": "signed-centimetres-int16",
        "mask_codec": "bitset-u8-v1",
    }
    identity_sha256 = _canonical_sha256(identity_document)
    pack_root = runs.parent / "polygon-review-packs" / run_id / f"review-{identity_sha256}"
    manifest_path = pack_root / "manifest.json"
    if manifest_path.is_file():
        return _read_completed_manifest(manifest_path, run_id, identity_sha256)

    frame_root = pack_root / "frames"
    frame_root.mkdir(parents=True, exist_ok=True)
    pending: list[tuple[int, str, str, str, Path]] = []
    completed: dict[str, dict[str, Any]] = {}
    for sequence, (frame_id, point_member, label_member) in enumerate(members):
        target = frame_root / f"{frame_id}.npz"
        metadata = _read_review_frame_metadata(target, frame_id, preview_points)
        if metadata is None:
            pending.append((sequence, frame_id, point_member, label_member, target))
        else:
            completed[frame_id] = metadata

    if pending:
        with ProcessPoolExecutor(
            max_workers=parallel_workers,
            initializer=_initialize_patchwork_worker,
            initargs=(patchwork_module_name,),
        ) as executor:
            futures = {
                executor.submit(
                    _build_review_frame_worker,
                    str(archive_path),
                    label_mapping,
                    frame_id,
                    point_member,
                    label_member,
                    str(target),
                    preview_points,
                ): frame_id
                for _, frame_id, point_member, label_member, target in pending
            }
            for future in as_completed(futures):
                frame_id = futures[future]
                completed[frame_id] = future.result()

    frames: list[dict[str, Any]] = []
    for sequence, (frame_id, _, _) in enumerate(members):
        source = report_frames[frame_id]
        nominal = source.get("nominal")
        if not isinstance(nominal, dict):
            raise GooseAdmissionError("qualification frame has no nominal metrics")
        current = _review_metrics(nominal.get("current"))
        patchwork = _review_metrics(nominal.get("patchworkpp"))
        item = completed[frame_id]
        frames.append(
            {
                "sequence": sequence,
                "frame_id": frame_id,
                "source_point_count": item["source_point_count"],
                "point_count": item["point_count"],
                "relative_path": f"frames/{frame_id}.npz",
                "sha256": item["sha256"],
                "byte_length": item["byte_length"],
                "current": current,
                "patchworkpp": patchwork,
                "ground_iou_delta": (patchwork["ground_iou"] - current["ground_iou"]),
            }
        )

    manifest = {
        **identity_document,
        "identity_sha256": identity_sha256,
        "source_id": GOOSE_SOURCE_ID,
        "frame_count": len(frames),
        "frames": frames,
        "safety": {
            "visualization_only": True,
            "actuator_authority": False,
            "navigation_or_safety_accepted": False,
        },
    }
    _write_json_once(manifest_path, manifest)
    return _read_completed_manifest(manifest_path, run_id, identity_sha256)


def _initialize_patchwork_worker(module_name: str) -> None:
    global _WORKER_PATCHWORK
    _WORKER_PATCHWORK = PatchworkPPGroundSegmenter.load(
        DEFAULT_GOOSE_PATCHWORK_PROFILE,
        module_name=module_name,
    )


def _build_review_frame_worker(
    archive_path: str,
    label_mapping: dict[int, dict[str, Any]],
    frame_id: str,
    point_member: str,
    label_member: str,
    target: str,
    preview_points: int,
) -> dict[str, Any]:
    if _WORKER_PATCHWORK is None:
        raise GooseAdmissionError("Patchwork++ review worker was not initialized")
    frame = _read_archive_frame(Path(archive_path), point_member, label_member)
    categories = _challenge_categories(frame, label_mapping)
    evaluated = categories != 0
    ground_truth = (categories == 2) | (categories == 3)
    xyzi = np.column_stack((frame.points_xyz_m, frame.remission)).astype(
        np.float32,
        copy=False,
    )
    current = LocalPercentileGroundSegmenter(DEFAULT_GROUND_BENCHMARK_PROFILE).segment(xyzi)
    patchwork = _WORKER_PATCHWORK.segment(xyzi)
    if current.ground_mask.shape != (frame.point_count,) or patchwork.ground_mask.shape != (
        frame.point_count,
    ):
        raise GooseAdmissionError("review provider result is not point-aligned")
    sample_count = min(frame.point_count, preview_points)
    indices = np.linspace(0, frame.point_count - 1, sample_count, dtype=np.int64)
    points_cm = np.rint(frame.points_xyz_m[indices] * 100.0)
    if not np.all(np.isfinite(points_cm)) or np.any(np.abs(points_cm) > 32_767):
        raise GooseAdmissionError("review point escaped the signed-centimetre envelope")
    flags = (
        evaluated[indices].astype(np.uint8)
        | (ground_truth[indices].astype(np.uint8) << 1)
        | (current.ground_mask[indices].astype(np.uint8) << 2)
        | (patchwork.ground_mask[indices].astype(np.uint8) << 3)
    )
    remission = np.asarray(frame.remission[indices], dtype=np.float32)
    remission_max = float(np.max(remission)) if remission.size else 0.0
    if remission_max <= 1.0:
        remission = remission * 255.0
    remission_u8 = np.clip(np.rint(remission), 0, 255).astype(np.uint8)
    path = Path(target)
    path.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile(dir=path.parent, delete=False) as temporary:
        temporary_path = Path(temporary.name)
        np.savez_compressed(
            temporary,
            schema=np.asarray([GOOSE_REVIEW_FRAME_SCHEMA]),
            frame_id=np.asarray([frame_id]),
            source_point_count=np.asarray([frame.point_count], dtype=np.int32),
            xyz_cm=np.ascontiguousarray(points_cm, dtype="<i2"),
            remission_u8=remission_u8,
            flags=flags,
        )
        temporary.flush()
        os.fsync(temporary.fileno())
    os.replace(temporary_path, path)
    return {
        "source_point_count": frame.point_count,
        "point_count": sample_count,
        "sha256": _sha256_file(path),
        "byte_length": path.stat().st_size,
    }


def _read_review_frame_metadata(
    path: Path,
    frame_id: str,
    maximum_points: int,
) -> dict[str, Any] | None:
    if not path.is_file():
        return None
    try:
        with np.load(path, allow_pickle=False) as source:
            schema = str(source["schema"][0])
            stored_frame_id = str(source["frame_id"][0])
            source_point_count = int(source["source_point_count"][0])
            xyz_cm = source["xyz_cm"]
            remission = source["remission_u8"]
            flags = source["flags"]
            point_count = int(xyz_cm.shape[0])
            if (
                schema != GOOSE_REVIEW_FRAME_SCHEMA
                or stored_frame_id != frame_id
                or xyz_cm.dtype != np.dtype("<i2")
                or xyz_cm.shape != (point_count, 3)
                or remission.dtype != np.uint8
                or remission.shape != (point_count,)
                or flags.dtype != np.uint8
                or flags.shape != (point_count,)
                or not 1 <= point_count <= maximum_points
                or source_point_count < point_count
            ):
                return None
    except (OSError, ValueError, KeyError, IndexError):
        return None
    return {
        "source_point_count": source_point_count,
        "point_count": point_count,
        "sha256": _sha256_file(path),
        "byte_length": path.stat().st_size,
    }


def _review_metrics(value: object) -> dict[str, float]:
    if not isinstance(value, dict) or not isinstance(value.get("metrics"), dict):
        raise GooseAdmissionError("qualification frame metrics are unavailable")
    metrics = value["metrics"]
    return {
        "ground_iou": _finite_fraction(metrics.get("ground_iou")),
        "natural_ground_recall": _finite_fraction(metrics.get("natural_ground_recall")),
        "obstacle_non_ground_recall": _finite_fraction(metrics.get("obstacle_non_ground_recall")),
        "latency_ms": _finite_nonnegative(value.get("latency_ms")),
    }


def _finite_fraction(value: object) -> float:
    number = _finite_nonnegative(value)
    if number > 1:
        raise GooseAdmissionError("qualification metric is outside 0..1")
    return number


def _finite_nonnegative(value: object) -> float:
    if not isinstance(value, (int, float)) or not math.isfinite(value) or value < 0:
        raise GooseAdmissionError("qualification metric is invalid")
    return float(value)


def _read_completed_manifest(
    path: Path,
    run_id: str,
    identity_sha256: str,
) -> dict[str, Any]:
    try:
        value = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
        raise GooseAdmissionError("GOOSE review manifest is invalid") from exc
    if (
        not isinstance(value, dict)
        or value.get("schema_version") != GOOSE_REVIEW_PACK_SCHEMA
        or value.get("source_run_id") != run_id
        or value.get("identity_sha256") != identity_sha256
        or not isinstance(value.get("frames"), list)
        or value.get("frame_count") != len(value["frames"])
    ):
        raise GooseAdmissionError("GOOSE review manifest contract differs")
    return value


def _canonical_sha256(value: dict[str, Any]) -> str:
    return hashlib.sha256(
        json.dumps(value, sort_keys=True, separators=(",", ":")).encode()
    ).hexdigest()


def _write_json_once(path: Path, value: dict[str, Any]) -> None:
    encoded = json.dumps(value, sort_keys=True, separators=(",", ":")).encode() + b"\n"
    if path.exists():
        if path.is_file() and path.read_bytes() == encoded:
            return
        raise GooseAdmissionError("immutable GOOSE review manifest already exists")
    path.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile(dir=path.parent, delete=False) as temporary:
        temporary_path = Path(temporary.name)
        temporary.write(encoded)
        temporary.flush()
        os.fsync(temporary.fileno())
    os.replace(temporary_path, path)
