"""Independent-label ground benchmark for an admitted GOOSE native scan."""

from __future__ import annotations

import hashlib
import json
import os
import tempfile
import time
from pathlib import Path
from typing import Any, Final

import numpy as np

from k1link.datasets.gateway import read_dataset_admission_manifest, read_semantic_kitti_frame
from k1link.datasets.goose_admission import (
    GOOSE_SOURCE_ID,
    GooseAdmissionError,
    read_goose_label_mapping,
)
from k1link.datasets.goose_profile import (
    DEFAULT_GOOSE_PATCHWORK_PROFILE,
    GoosePatchworkProfile,
)
from k1link.ground_segmentation import (
    DEFAULT_GROUND_BENCHMARK_PROFILE,
    GroundBenchmarkProfile,
    GroundSegmenter,
    LocalPercentileGroundSegmenter,
    PatchworkPPGroundSegmenter,
)

GOOSE_GROUND_BENCHMARK_SCHEMA: Final = "missioncore.goose-ground-benchmark/v1"
GOOSE_GROUND_PREVIEW_SCHEMA: Final = "missioncore.dataset-ground-comparison-preview/v1"
GOOSE_GROUND_AB_BENCHMARK_SCHEMA: Final = "missioncore.goose-ground-ab-benchmark/v1"
GOOSE_GROUND_AB_PREVIEW_SCHEMA: Final = "missioncore.dataset-ground-comparison-preview/v2"
MAX_PREVIEW_POINTS: Final = 50_000


def benchmark_goose_current_ground(
    dataset_root: Path,
    *,
    profile: GroundBenchmarkProfile = DEFAULT_GROUND_BENCHMARK_PROFILE,
    preview_points: int = MAX_PREVIEW_POINTS,
) -> dict[str, Any]:
    """Score the existing Mission Core ground baseline on one admitted frame."""

    root = dataset_root.expanduser().absolute()
    if not _is_worker_dataset_root(root):
        raise GooseAdmissionError("GOOSE benchmark requires the canonical worker D root")
    manifest = read_dataset_admission_manifest(root / "state/goose-3d-v2025-08-22.json")
    if manifest["status"] != "frame-ready":
        raise GooseAdmissionError("GOOSE frame is not admitted for benchmarking")
    archive_sha256 = str(manifest["archive"]["sha256"])
    frame_id = str(manifest["frame"]["frame_id"])
    install = root / "goose-3d/v2025-08-22/installs" / archive_sha256
    frame_root = install / "frames" / frame_id
    frame = read_semantic_kitti_frame(
        frame_root / "points.bin",
        frame_root / "labels.label",
    )
    labels = read_goose_label_mapping(install / "goose_label_mapping.csv")
    categories = np.asarray(
        [labels[int(value)]["challenge_category_id"] for value in frame.semantic_labels],
        dtype=np.uint8,
    )
    evaluated = categories != 0
    ground_truth = (categories == 2) | (categories == 3)
    xyzi = np.column_stack((frame.points_xyz_m, frame.remission)).astype(
        np.float32,
        copy=False,
    )
    segmenter = LocalPercentileGroundSegmenter(profile)
    started = time.perf_counter_ns()
    prediction = segmenter.segment(xyzi)
    wall_latency_ms = (time.perf_counter_ns() - started) / 1_000_000
    metrics = _binary_metrics(
        prediction.ground_mask,
        ground_truth,
        evaluated,
    )
    metrics.update(
        {
            "artificial_ground_recall": _recall(
                prediction.ground_mask,
                categories == 2,
            ),
            "natural_ground_recall": _recall(
                prediction.ground_mask,
                categories == 3,
            ),
            "obstacle_non_ground_recall": _recall(
                ~prediction.ground_mask,
                categories == 4,
            ),
        }
    )
    profile_document = profile.to_dict()
    identity_document = {
        "source_id": GOOSE_SOURCE_ID,
        "archive_sha256": archive_sha256,
        "frame_id": frame_id,
        "source_point_count": frame.point_count,
        "profile": profile_document,
        "provider": dict(segmenter.identity),
    }
    identity_sha256 = _canonical_sha256(identity_document)
    output_root = install / "benchmarks" / f"current-ground-{identity_sha256}"
    output_root.mkdir(parents=True, exist_ok=True)
    report_path = output_root / "report.json"
    if report_path.is_file():
        try:
            existing = json.loads(report_path.read_text(encoding="utf-8"))
        except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
            raise GooseAdmissionError("existing GOOSE benchmark report is invalid") from exc
        if (
            not isinstance(existing, dict)
            or existing.get("schema_version") != GOOSE_GROUND_BENCHMARK_SCHEMA
            or existing.get("identity_sha256") != identity_sha256
            or not (output_root / "prediction.npz").is_file()
            or not (output_root / "preview.json").is_file()
        ):
            raise GooseAdmissionError("existing GOOSE benchmark is incomplete")
        return existing
    prediction_path = output_root / "prediction.npz"
    _atomic_npz(
        prediction_path,
        current_ground=prediction.ground_mask.astype(np.uint8),
        ground_truth_ground=ground_truth.astype(np.uint8),
        evaluated=evaluated.astype(np.uint8),
    )
    prediction_sha256 = _sha256_file(prediction_path)

    sample_count = min(frame.point_count, preview_points)
    indices = np.linspace(0, frame.point_count - 1, sample_count, dtype=np.int64)
    preview = {
        "schema_version": GOOSE_GROUND_PREVIEW_SCHEMA,
        "source_id": GOOSE_SOURCE_ID,
        "frame_id": frame_id,
        "sampling": "deterministic-even-index",
        "point_count": sample_count,
        "current_ground": prediction.ground_mask[indices].astype(np.uint8).tolist(),
        "ground_truth_ground": ground_truth[indices].astype(np.uint8).tolist(),
        "evaluated": evaluated[indices].astype(np.uint8).tolist(),
        "disagreement": (
            evaluated[indices] & (prediction.ground_mask[indices] != ground_truth[indices])
        )
        .astype(np.uint8)
        .tolist(),
        "metrics": metrics,
        "latency_ms": prediction.latency_ms,
        "provider": dict(segmenter.identity),
        "safety": {
            "qualification_only": True,
            "navigation_or_safety_accepted": False,
        },
    }
    preview_path = output_root / "preview.json"
    _atomic_json(preview_path, preview)
    report = {
        "schema_version": GOOSE_GROUND_BENCHMARK_SCHEMA,
        "identity_sha256": identity_sha256,
        "identity": identity_document,
        "ground_truth": {
            "source": "GOOSE point-wise semantic labels",
            "ground_categories": ["artificial_ground", "natural_ground"],
            "void_excluded": True,
            "evaluated_points": int(np.count_nonzero(evaluated)),
        },
        "metrics": metrics,
        "latency": {
            "provider_ms": prediction.latency_ms,
            "wall_ms": wall_latency_ms,
        },
        "artifacts": {
            "prediction_sha256": prediction_sha256,
            "preview_sha256": _sha256_file(preview_path),
        },
        "decision": {
            "status": "diagnostic-baseline",
            "promoted": False,
            "reason": "one public frame does not qualify a production ground provider",
        },
    }
    _atomic_json(report_path, report)
    return report


def benchmark_goose_patchwork_ground(
    dataset_root: Path,
    *,
    current_profile: GroundBenchmarkProfile = DEFAULT_GROUND_BENCHMARK_PROFILE,
    patchwork_profile: GoosePatchworkProfile = DEFAULT_GOOSE_PATCHWORK_PROFILE,
    patchwork: GroundSegmenter | None = None,
    patchwork_module_name: str = "pypatchworkpp",
    preview_points: int = MAX_PREVIEW_POINTS,
) -> dict[str, Any]:
    """Score current and pinned Patchwork++ providers against one GOOSE frame."""

    root = dataset_root.expanduser().absolute()
    if not _is_worker_dataset_root(root):
        raise GooseAdmissionError("GOOSE benchmark requires the canonical worker D root")
    manifest = read_dataset_admission_manifest(root / "state/goose-3d-v2025-08-22.json")
    if manifest["status"] != "frame-ready":
        raise GooseAdmissionError("GOOSE frame is not admitted for benchmarking")
    archive_sha256 = str(manifest["archive"]["sha256"])
    frame_id = str(manifest["frame"]["frame_id"])
    install = root / "goose-3d/v2025-08-22/installs" / archive_sha256
    frame_root = install / "frames" / frame_id
    frame = read_semantic_kitti_frame(
        frame_root / "points.bin",
        frame_root / "labels.label",
    )
    labels = read_goose_label_mapping(install / "goose_label_mapping.csv")
    categories = np.asarray(
        [labels[int(value)]["challenge_category_id"] for value in frame.semantic_labels],
        dtype=np.uint8,
    )
    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(current_profile)
    candidate = (
        patchwork
        if patchwork is not None
        else PatchworkPPGroundSegmenter.load(
            patchwork_profile,
            module_name=patchwork_module_name,
        )
    )
    current_result = current.segment(xyzi)
    candidate_result = candidate.segment(xyzi)
    for result, label in (
        (current_result, "current"),
        (candidate_result, "Patchwork++"),
    ):
        if result.ground_mask.shape != (frame.point_count,) or result.assigned_mask.shape != (
            frame.point_count,
        ):
            raise GooseAdmissionError(f"{label} result is not point-aligned")
    current_metrics = _ground_metrics(
        current_result.ground_mask,
        ground_truth,
        evaluated,
        categories,
    )
    candidate_metrics = _ground_metrics(
        candidate_result.ground_mask,
        ground_truth,
        evaluated,
        categories,
    )
    candidate_assigned_fraction = float(np.mean(candidate_result.assigned_mask))
    profile_document = patchwork_profile.to_dict()
    current_profile_document = _goose_current_profile_document(current_profile)
    identity_document = {
        "source_id": GOOSE_SOURCE_ID,
        "archive_sha256": archive_sha256,
        "frame_id": frame_id,
        "source_point_count": frame.point_count,
        "input_profile": profile_document,
        "current_profile": current_profile_document,
        "providers": {
            "current": dict(current.identity),
            "patchworkpp": dict(candidate.identity),
        },
    }
    identity_sha256 = _canonical_sha256(identity_document)
    output_root = install / "benchmarks" / f"ground-ab-{identity_sha256}"
    output_root.mkdir(parents=True, exist_ok=True)
    report_path = output_root / "report.json"
    if report_path.is_file():
        try:
            existing = json.loads(report_path.read_text(encoding="utf-8"))
        except (OSError, UnicodeDecodeError, json.JSONDecodeError) as exc:
            raise GooseAdmissionError("existing GOOSE A/B report is invalid") from exc
        if (
            not isinstance(existing, dict)
            or existing.get("schema_version") != GOOSE_GROUND_AB_BENCHMARK_SCHEMA
            or existing.get("identity_sha256") != identity_sha256
            or not (output_root / "prediction.npz").is_file()
            or not (output_root / "preview.json").is_file()
        ):
            raise GooseAdmissionError("existing GOOSE A/B benchmark is incomplete")
        return existing

    prediction_path = output_root / "prediction.npz"
    _atomic_npz(
        prediction_path,
        current_ground=current_result.ground_mask.astype(np.uint8),
        patchwork_ground=candidate_result.ground_mask.astype(np.uint8),
        patchwork_assigned=candidate_result.assigned_mask.astype(np.uint8),
        ground_truth_ground=ground_truth.astype(np.uint8),
        evaluated=evaluated.astype(np.uint8),
    )
    sample_count = min(frame.point_count, preview_points)
    indices = np.linspace(0, frame.point_count - 1, sample_count, dtype=np.int64)
    current_disagreement = evaluated & (current_result.ground_mask != ground_truth)
    candidate_disagreement = evaluated & (candidate_result.ground_mask != ground_truth)
    preview = {
        "schema_version": GOOSE_GROUND_AB_PREVIEW_SCHEMA,
        "source_id": GOOSE_SOURCE_ID,
        "frame_id": frame_id,
        "sampling": "deterministic-even-index",
        "point_count": sample_count,
        "current_ground": current_result.ground_mask[indices].astype(np.uint8).tolist(),
        "patchwork_ground": candidate_result.ground_mask[indices].astype(np.uint8).tolist(),
        "patchwork_assigned": (candidate_result.assigned_mask[indices].astype(np.uint8).tolist()),
        "ground_truth_ground": ground_truth[indices].astype(np.uint8).tolist(),
        "evaluated": evaluated[indices].astype(np.uint8).tolist(),
        "current_disagreement": (current_disagreement[indices].astype(np.uint8).tolist()),
        "patchwork_disagreement": (candidate_disagreement[indices].astype(np.uint8).tolist()),
        "metrics": {
            "current": current_metrics,
            "patchworkpp": candidate_metrics,
        },
        "latency_ms": {
            "current": current_result.latency_ms,
            "patchworkpp": candidate_result.latency_ms,
        },
        "patchwork_assigned_fraction": candidate_assigned_fraction,
        "providers": {
            "current": dict(current.identity),
            "patchworkpp": dict(candidate.identity),
        },
        "input_profile": profile_document,
        "safety": {
            "qualification_only": True,
            "navigation_or_safety_accepted": False,
        },
    }
    preview_path = output_root / "preview.json"
    _atomic_json(preview_path, preview)
    winner = (
        "patchworkpp"
        if candidate_metrics["ground_iou"] > current_metrics["ground_iou"]
        else "current"
        if current_metrics["ground_iou"] > candidate_metrics["ground_iou"]
        else "tie"
    )
    report = {
        "schema_version": GOOSE_GROUND_AB_BENCHMARK_SCHEMA,
        "identity_sha256": identity_sha256,
        "identity": identity_document,
        "ground_truth": {
            "source": "GOOSE point-wise semantic labels",
            "ground_categories": ["artificial_ground", "natural_ground"],
            "void_excluded": True,
            "evaluated_points": int(np.count_nonzero(evaluated)),
        },
        "providers": {
            "current": {
                "metrics": current_metrics,
                "latency_ms": current_result.latency_ms,
                "assigned_fraction": 1.0,
            },
            "patchworkpp": {
                "metrics": candidate_metrics,
                "latency_ms": candidate_result.latency_ms,
                "assigned_fraction": candidate_assigned_fraction,
            },
        },
        "artifacts": {
            "prediction_sha256": _sha256_file(prediction_path),
            "preview_sha256": _sha256_file(preview_path),
        },
        "decision": {
            "status": "one-frame-diagnostic",
            "promoted": False,
            "winner_by_ground_iou": winner,
            "reason": "one public frame does not qualify a production ground provider",
            "next_gate": "validation-split qualification with degradation profiles",
        },
        "scope": {
            "patchworkpp_input_admitted": True,
            "normalized_scan_produced": False,
            "complete_vehicle_transform_known": False,
        },
    }
    _atomic_json(report_path, report)
    return report


def _is_worker_dataset_root(root: Path) -> bool:
    normalized = str(root).replace("\\", "/").rstrip("/").lower()
    return normalized == "/mnt/d/ndc_missioncore/datasets"


def _binary_metrics(
    prediction: np.ndarray[Any, Any],
    ground_truth: np.ndarray[Any, Any],
    evaluated: np.ndarray[Any, Any],
) -> dict[str, float | int]:
    predicted = prediction & evaluated
    truth = ground_truth & evaluated
    true_positive = int(np.count_nonzero(predicted & truth))
    false_positive = int(np.count_nonzero(predicted & ~truth & evaluated))
    false_negative = int(np.count_nonzero(~predicted & truth))
    true_negative = int(np.count_nonzero(~predicted & ~truth & evaluated))
    precision = _ratio(true_positive, true_positive + false_positive)
    recall = _ratio(true_positive, true_positive + false_negative)
    return {
        "true_positive": true_positive,
        "false_positive": false_positive,
        "false_negative": false_negative,
        "true_negative": true_negative,
        "precision": precision,
        "recall": recall,
        "f1": _ratio(2 * precision * recall, precision + recall),
        "ground_iou": _ratio(
            true_positive,
            true_positive + false_positive + false_negative,
        ),
        "accuracy": _ratio(
            true_positive + true_negative,
            true_positive + true_negative + false_positive + false_negative,
        ),
    }


def _ground_metrics(
    prediction: np.ndarray[Any, Any],
    ground_truth: np.ndarray[Any, Any],
    evaluated: np.ndarray[Any, Any],
    categories: np.ndarray[Any, Any],
) -> dict[str, float | int]:
    metrics = _binary_metrics(prediction, ground_truth, evaluated)
    metrics.update(
        {
            "artificial_ground_recall": _recall(prediction, categories == 2),
            "natural_ground_recall": _recall(prediction, categories == 3),
            "obstacle_non_ground_recall": _recall(~prediction, categories == 4),
        }
    )
    return metrics


def _goose_current_profile_document(
    profile: GroundBenchmarkProfile,
) -> dict[str, object]:
    parameters = profile.to_dict()["current_baseline"]
    return {
        "schema_version": "missioncore.goose-current-ground-profile/v1",
        "profile_id": "goose-native-local-percentile/v1",
        "source_id": GOOSE_SOURCE_ID,
        "representation": "native-scan",
        "sensor_frame": "sensor/lidar/vls128_roof",
        "provider": parameters,
        "scope": {
            "complete_vehicle_transform_required": False,
            "normalized_scan_produced": False,
        },
        "authority": {
            "qualification_only": True,
            "navigation_or_safety_accepted": False,
        },
    }


def _recall(prediction: np.ndarray[Any, Any], truth: np.ndarray[Any, Any]) -> float:
    return _ratio(
        int(np.count_nonzero(prediction & truth)),
        int(np.count_nonzero(truth)),
    )


def _ratio(numerator: float | int, denominator: float | int) -> float:
    return float(numerator / denominator) if denominator else 0.0


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


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


def _atomic_npz(path: Path, **arrays: np.ndarray[Any, Any]) -> None:
    if path.exists():
        raise GooseAdmissionError("immutable GOOSE benchmark artifact already exists")
    with tempfile.NamedTemporaryFile(dir=path.parent, suffix=".npz", delete=False) as temporary:
        temporary_path = Path(temporary.name)
    try:
        np.savez_compressed(temporary_path, **arrays)  # type: ignore[arg-type]
        with temporary_path.open("rb") as source:
            os.fsync(source.fileno())
        os.replace(temporary_path, path)
    except Exception:
        temporary_path.unlink(missing_ok=True)
        raise


def _atomic_json(path: Path, value: dict[str, Any]) -> None:
    if path.exists():
        raise GooseAdmissionError("immutable GOOSE benchmark artifact already exists")
    encoded = json.dumps(value, sort_keys=True, separators=(",", ":")).encode("utf-8") + b"\n"
    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)
