from __future__ import annotations

import json
import struct
import zipfile
from collections.abc import Mapping
from pathlib import Path

import numpy as np
import numpy.typing as npt
import pytest

from k1link.datasets import goose_qualification
from k1link.datasets.goose_qualification import (
    DegradationProfile,
    qualify_goose_ground,
)
from k1link.ground_segmentation import GroundSegmentation
from k1link.simulation import QualificationRunStore, RunState


class _NeverGround:
    @property
    def identity(self) -> Mapping[str, object]:
        return {"provider_id": "test-current", "revision": "1"}

    def segment(self, xyzi: npt.NDArray[np.float32]) -> GroundSegmentation:
        count = xyzi.shape[0]
        return GroundSegmentation(
            ground_mask=np.zeros(count, dtype=np.bool_),
            assigned_mask=np.ones(count, dtype=np.bool_),
            latency_ms=2.0,
        )


class _NegativeZGround:
    @property
    def identity(self) -> Mapping[str, object]:
        return {"provider_id": "test-patchwork", "revision": "1"}

    def segment(self, xyzi: npt.NDArray[np.float32]) -> GroundSegmentation:
        count = xyzi.shape[0]
        return GroundSegmentation(
            ground_mask=np.asarray(xyzi[:, 2] < 0, dtype=np.bool_),
            assigned_mask=np.ones(count, dtype=np.bool_),
            latency_ms=3.0,
        )


def _frame_bytes(offset: float) -> tuple[bytes, bytes]:
    points = (
        (-1.0 + offset, 0.0, -2.0, 0.1),
        (0.0 + offset, 0.0, -2.0, 0.2),
        (1.0 + offset, 0.0, 2.0, 0.3),
        (2.0 + offset, 0.0, 2.0, 0.4),
    )
    labels = (1, 3, 2, 2)
    return (
        b"".join(struct.pack("<ffff", *point) for point in points),
        b"".join(struct.pack("<I", label) for label in labels),
    )


def _dataset(tmp_path: Path) -> Path:
    root = tmp_path / "datasets"
    archive_sha = "a" * 64
    archive = root / "goose-3d/v2025-08-22/archives/goose_3d_val.zip"
    archive.parent.mkdir(parents=True)
    with zipfile.ZipFile(archive, "w") as target:
        for index in range(2):
            frame_id = f"2022-01-01_flight__0001_{index:019d}"
            points, labels = _frame_bytes(index * 0.1)
            target.writestr(f"goose/lidar/val/{frame_id}_vls128.bin", points)
            target.writestr(f"goose/labels/val/{frame_id}_goose.label", labels)
    state = root / "state/goose-3d-v2025-08-22.json"
    state.parent.mkdir(parents=True)
    state.write_text(json.dumps({"archive": {"sha256": archive_sha}}), encoding="utf-8")
    install = root / "goose-3d/v2025-08-22/installs" / archive_sha
    install.mkdir(parents=True)
    (install / "goose_label_mapping.csv").write_text(
        "class_name,label_key,hex\nasphalt,1,#111111\nrock,2,#222222\nsoil,3,#333333\n",
        encoding="utf-8",
    )
    return root


def test_full_split_qualification_is_sealed_and_resumable(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    root = _dataset(tmp_path)
    runs_root = tmp_path / "runs"
    monkeypatch.setattr(goose_qualification, "_is_worker_dataset_root", lambda _: True)

    report = qualify_goose_ground(
        root,
        runs_root,
        mission_core_commit="1" * 40,
        parallel_workers=1,
        expected_frame_count=2,
        current_segmenter=_NeverGround(),
        patchwork_segmenter=_NegativeZGround(),
        degradations=(DegradationProfile("range-20m", "maximum-range-m", 20.0),),
    )

    assert report["frame_count"] == 2
    assert report["decision"]["status"] == "shadow-candidate"
    assert report["aggregates"]["current"]["micro"]["ground_iou"] == 0
    assert report["aggregates"]["patchworkpp"]["micro"]["ground_iou"] == 1
    assert len(report["frames"]) == 2
    run = QualificationRunStore(runs_root, read_only=True).load(report["run_id"])
    assert run.state is RunState.COMPLETED
    assert {artifact.kind for artifact in run.artifacts} == {
        "goose-ground-qualification-report",
        "goose-ground-qualification-failure-preview",
    }

    resumed = qualify_goose_ground(
        root,
        runs_root,
        mission_core_commit="1" * 40,
        parallel_workers=1,
        expected_frame_count=2,
        current_segmenter=_NeverGround(),
        patchwork_segmenter=_NegativeZGround(),
        degradations=(DegradationProfile("range-20m", "maximum-range-m", 20.0),),
    )
    assert resumed["identity_sha256"] == report["identity_sha256"]