from __future__ import annotations

from dataclasses import replace

import numpy as np
import pytest

from k1link.compute.sensor_representation import K1_LIO_PCL_CAPABILITIES
from k1link.compute.track_geometry import (
    PointSlab,
    TrackGeometry,
    TrackGeometryContractError,
    TrackGeometryCurrentness,
    TrackGeometryEvidenceState,
    TrackGeometryFrame,
    TrackGeometryMetricBasis,
    TrackGeometryOwnerKind,
    TrackGeometrySourceBinding,
)


def _binding() -> TrackGeometrySourceBinding:
    return TrackGeometrySourceBinding(
        source_pack_id="e10-lidar-pack-" + "a" * 64,
        source_session_id="source-session",
        representation_profile_id=K1_LIO_PCL_CAPABILITIES.profile_id,
        e31_qualification_id="e31-source-qualification-" + "b" * 64,
        calibration_sha256="c" * 64,
        coordinate_frame="map",
        time_basis="recorded-host-arrival-best-effort",
        selected_offset_ms=0,
    )


def _agree(owner_key: str = "track:7") -> TrackGeometry:
    return TrackGeometry(
        owner_key=owner_key,
        owner_kind=TrackGeometryOwnerKind.CAMERA_TRACK,
        evidence_state=TrackGeometryEvidenceState.AGREE,
        currentness=TrackGeometryCurrentness.CURRENT,
        metric_basis=TrackGeometryMetricBasis.CURRENT_POINTS,
        reason_codes=("connected-occupied-support",),
        semantic_track_id=7,
        semantic_label="car",
        bbox_xyxy=(10.0, 20.0, 50.0, 60.0),
        range_m=8.2,
    )


def _camera_only() -> TrackGeometry:
    return TrackGeometry(
        owner_key="track:8",
        owner_kind=TrackGeometryOwnerKind.CAMERA_TRACK,
        evidence_state=TrackGeometryEvidenceState.CAMERA_ONLY,
        currentness=TrackGeometryCurrentness.CURRENT,
        metric_basis=TrackGeometryMetricBasis.UNAVAILABLE,
        reason_codes=("sparse-support",),
        semantic_track_id=8,
        semantic_label="person",
        bbox_xyxy=(60.0, 20.0, 80.0, 90.0),
    )


def _geometry_only() -> TrackGeometry:
    return TrackGeometry(
        owner_key="geometry:2",
        owner_kind=TrackGeometryOwnerKind.GEOMETRY_CLUSTER,
        evidence_state=TrackGeometryEvidenceState.GEOMETRY_ONLY,
        currentness=TrackGeometryCurrentness.CURRENT,
        metric_basis=TrackGeometryMetricBasis.CURRENT_POINTS,
        reason_codes=("unassociated-occupied-component",),
        range_m=6.4,
    )


def _point_slab() -> PointSlab:
    return PointSlab(
        frame_index=10,
        source_frame_index=1010,
        source_point_count=100,
        coordinate_frame="map",
        owner_keys=("track:7", "geometry:2"),
        source_indices=np.asarray([1, 2, 50], dtype=np.int64),
        points_xyz_m=np.asarray(
            [
                [1.0, 0.0, 2.0],
                [1.1, 0.0, 2.0],
                [5.0, 1.0, 0.5],
            ],
            dtype=np.float32,
        ),
        owner_indices=np.asarray([0, 0, 1], dtype=np.uint32),
    )


def test_track_geometry_frame_round_trips_exact_source_ownership() -> None:
    frame = TrackGeometryFrame(
        binding=_binding(),
        frame_index=10,
        source_frame_index=1010,
        session_seconds=12.5,
        source_available=True,
        point_slab=_point_slab(),
        geometries=(_agree(), _camera_only(), _geometry_only()),
    )

    document = frame.to_dict()
    restored = TrackGeometryFrame.from_dict(document)

    assert restored.to_dict() == document
    assert restored.point_slab.owned_source_indices("track:7").tolist() == [1, 2]
    assert restored.point_slab.owned_source_indices("geometry:2").tolist() == [50]
    assert restored.point_slab.source_indices.flags.writeable is False
    assert document["policy"] == {
        "camera_owns_semantics": True,
        "one_owner_per_source_point": True,
        "current_held_persistent_are_separate": True,
        "absence_of_points_means_free": False,
        "unknown_remains_unknown": True,
    }


def test_point_slab_rejects_duplicate_or_out_of_frame_source_indices() -> None:
    slab = _point_slab()
    with pytest.raises(
        TrackGeometryContractError,
        match="unique and frame-local",
    ):
        replace(
            slab,
            source_indices=np.asarray([1, 1, 50], dtype=np.int64),
            owner_indices=np.asarray([0, 1, 1], dtype=np.uint32),
        )
    with pytest.raises(
        TrackGeometryContractError,
        match="unique and frame-local",
    ):
        replace(
            slab,
            source_indices=np.asarray([1, 2, 100], dtype=np.int64),
        )


def test_point_slab_rejects_unknown_or_multiple_owners() -> None:
    slab = _point_slab()
    with pytest.raises(TrackGeometryContractError, match="owner keys must be unique"):
        replace(slab, owner_keys=("track:7", "track:7"))
    with pytest.raises(TrackGeometryContractError, match="out of range"):
        replace(
            slab,
            owner_indices=np.asarray([0, 0, 2], dtype=np.uint32),
        )
    with pytest.raises(TrackGeometryContractError, match="unused owners"):
        replace(
            slab,
            owner_indices=np.asarray([0, 0, 0], dtype=np.uint32),
        )


def test_frame_rejects_current_points_claimed_by_two_geometry_products() -> None:
    slab = _point_slab()
    with pytest.raises(
        TrackGeometryContractError,
        match="geometry owner keys must be unique",
    ):
        TrackGeometryFrame(
            binding=_binding(),
            frame_index=10,
            source_frame_index=1010,
            session_seconds=12.5,
            source_available=True,
            point_slab=slab,
            geometries=(_agree(), _agree(), _geometry_only()),
        )


def test_camera_only_and_conflict_cannot_invent_metric_range() -> None:
    with pytest.raises(
        TrackGeometryContractError,
        match="unavailable metric geometry cannot publish range",
    ):
        replace(_camera_only(), range_m=3.0)
    with pytest.raises(
        TrackGeometryContractError,
        match="cannot invent metric geometry",
    ):
        replace(
            _camera_only(),
            evidence_state=TrackGeometryEvidenceState.CONFLICT,
            metric_basis=TrackGeometryMetricBasis.CURRENT_POINTS,
            range_m=3.0,
        )


def test_geometry_only_cannot_invent_camera_semantics() -> None:
    with pytest.raises(
        TrackGeometryContractError,
        match="cannot invent camera semantics",
    ):
        replace(
            _geometry_only(),
            semantic_track_id=9,
            semantic_label="car",
            bbox_xyxy=(1.0, 1.0, 2.0, 2.0),
        )


def test_held_geometry_keeps_provenance_but_owns_no_current_points() -> None:
    held = TrackGeometry(
        owner_key="track:7",
        owner_kind=TrackGeometryOwnerKind.CAMERA_TRACK,
        evidence_state=TrackGeometryEvidenceState.UNKNOWN,
        currentness=TrackGeometryCurrentness.HELD,
        metric_basis=TrackGeometryMetricBasis.HELD_LAST_CURRENT,
        reason_codes=("temporally-held-last-supported",),
        semantic_track_id=7,
        semantic_label="car",
        bbox_xyxy=(10.0, 20.0, 50.0, 60.0),
        range_m=8.2,
        held_from_frame_index=9,
    )
    empty_slab = PointSlab(
        frame_index=10,
        source_frame_index=1010,
        source_point_count=100,
        coordinate_frame="map",
        owner_keys=(),
        source_indices=np.empty(0, dtype=np.int64),
        points_xyz_m=np.empty((0, 3), dtype=np.float32),
        owner_indices=np.empty(0, dtype=np.uint32),
    )
    frame = TrackGeometryFrame(
        binding=_binding(),
        frame_index=10,
        source_frame_index=1010,
        session_seconds=12.5,
        source_available=True,
        point_slab=empty_slab,
        geometries=(held,),
    )

    assert TrackGeometryFrame.from_dict(frame.to_dict()).to_dict() == frame.to_dict()
    with pytest.raises(
        TrackGeometryContractError,
        match="current point ownership and metric basis disagree",
    ):
        TrackGeometryFrame(
            binding=_binding(),
            frame_index=10,
            source_frame_index=1010,
            session_seconds=12.5,
            source_available=True,
            point_slab=replace(
                _point_slab(),
                owner_keys=("track:7", "geometry:2"),
            ),
            geometries=(held, _geometry_only()),
        )


def test_persistent_geometry_is_a_separate_model_not_current_source_points() -> None:
    persistent = TrackGeometry(
        owner_key="persistent:wall:1",
        owner_kind=TrackGeometryOwnerKind.GEOMETRY_CLUSTER,
        evidence_state=TrackGeometryEvidenceState.GEOMETRY_ONLY,
        currentness=TrackGeometryCurrentness.PERSISTENT,
        metric_basis=TrackGeometryMetricBasis.PERSISTENT_MODEL,
        reason_codes=("persistent-occupied-model",),
        range_m=5.0,
        persistent_model_id="local-occupied-model/v1",
    )

    assert persistent.semantic_label is None
    with pytest.raises(
        TrackGeometryContractError,
        match="current-point metrics require current geometry",
    ):
        replace(
            persistent,
            metric_basis=TrackGeometryMetricBasis.CURRENT_POINTS,
        )


def test_unavailable_source_cannot_publish_current_point_geometry() -> None:
    with pytest.raises(
        TrackGeometryContractError,
        match="unavailable source cannot publish current metric points",
    ):
        TrackGeometryFrame(
            binding=_binding(),
            frame_index=10,
            source_frame_index=1010,
            session_seconds=12.5,
            source_available=False,
            point_slab=_point_slab(),
            geometries=(_agree(), _geometry_only()),
        )


def test_contract_rejects_authority_and_unknown_fields() -> None:
    frame = TrackGeometryFrame(
        binding=_binding(),
        frame_index=10,
        source_frame_index=1010,
        session_seconds=12.5,
        source_available=True,
        point_slab=_point_slab(),
        geometries=(_agree(), _camera_only(), _geometry_only()),
    )
    document = frame.to_dict()
    authority = document["authority"]
    assert isinstance(authority, dict)
    authority["navigation_or_safety_accepted"] = True
    with pytest.raises(TrackGeometryContractError, match="cannot grant"):
        TrackGeometryFrame.from_dict(document)

    document = frame.to_dict()
    document["free_space"] = True
    with pytest.raises(TrackGeometryContractError, match="fields are incompatible"):
        TrackGeometryFrame.from_dict(document)