from __future__ import annotations

import hashlib
from pathlib import Path
from types import SimpleNamespace

import numpy as np
import pytest
import rerun as rr
from fastapi import FastAPI
from fastapi.testclient import TestClient

from k1link.laboratory.canonical_rerun_overlay import CanonicalLabReplayArtifact
from k1link.observatory.portable_replay import PortableReplayService
from k1link.observatory.portable_tgs_replay import (
    PortableReplayError,
    TgsReplayData,
    log_tgs,
    verified_file,
    world_cells,
)
from k1link.web.portable_replay_api import build_portable_replay_router

RESULT = "m49-tgs-portable-review-" + "a" * 64
BASE = "b" * 64


@pytest.mark.parametrize(
    "damage", [None, "frame-time", "source-bundle", "nan-height", "state", "missing"]
)
def test_array_admission_checks_clocks_support_and_identity(tmp_path, monkeypatch, damage):
    from k1link.artifact_gateway import CentralArtifactStore
    from k1link.observatory.portable_result_contract import canonical_json
    from k1link.observatory.portable_tgs_replay import STATE_CODES, load_tgs_data

    store = CentralArtifactStore(tmp_path / "store", create=True)
    members = []

    def member(role, payload):
        path = tmp_path / role
        path.write_bytes(payload)
        item = store.publish_file(path)
        members.append({"role": role, "sha256": item.sha256, "byte_length": item.byte_length})
        return item.sha256

    source_row = {
        "timeline_frame_index": 0,
        "source_frame_index": 1,
        "session_seconds": 1.0,
        "sample_available": damage != "missing",
        "available_slot": 0,
        "position_map_m": [10.0, 20.0, 3.0],
    }
    monkeypatch.setattr(
        "k1link.observatory.portable_tgs_replay.read_m49_source_index",
        lambda *args, **kwargs: (source_row,),
    )
    member("source-stage-index", b"test")
    frame = dict(source_row)
    if damage == "frame-time":
        frame["session_seconds"] = 2.0
    member("frame-index", canonical_json(frame) + b"\n")
    identity = {
        "source": {"source_session_id": "source", "source_bundle_sha256": BASE},
        "profile": {"profile_sha256": "c" * 64},
    }
    stage_sha = hashlib.sha256(canonical_json(identity)).hexdigest()
    stage_manifest = member(
        "source-stage-manifest",
        canonical_json({"identity": identity, "identity_sha256": stage_sha}),
    )
    import io

    for role, array in (
        ("costmap-cell-centers", np.array([[1.0, 2.0]], dtype="<f4")),
        ("costmap-states", np.array([[4 if damage == "state" else 2]], dtype="u1")),
        (
            "costmap-z-bounds",
            np.array([[[np.nan if damage == "nan-height" else -1.0, 1.0]]], dtype="<f4"),
        ),
    ):
        buffer = io.BytesIO()
        np.save(buffer, array, allow_pickle=False)
        member(role, buffer.getvalue())
    view = {
        "result_id": RESULT,
        "source_session_id": "source",
        "artifacts": members,
        "result_document": {
            "schema_version": "missioncore.recorded-tgs-costmap-review/v2",
            "result_id": RESULT,
            "timeline": {"frame_count": 1},
            "costmap": {
                "cell_count": 1,
                "cell_size_m": 0.45,
                "coordinate_frame": "map-gravity-local",
                "state_codes": STATE_CODES,
            },
            "profile": identity["profile"],
            "source_stage": {"identity_sha256": stage_sha, "manifest_sha256": stage_manifest},
        },
    }
    if damage is None:
        assert load_tgs_data(view, store, source_bundle_sha256=BASE).states[0, 0] == 2
    else:
        with pytest.raises(PortableReplayError):
            load_tgs_data(
                view, store, source_bundle_sha256="d" * 64 if damage == "source-bundle" else BASE
            )


def data() -> TgsReplayData:
    return TgsReplayData(
        (
            {"session_seconds": 1.0, "sample_available": True, "position_map_m": [10, 20, 3]},
            {"session_seconds": 2.0, "sample_available": False, "position_map_m": None},
        ),
        np.array([[1.0, 2.0], [2.0, 3.0]]),
        np.array([[2, 0], [0, 0]], dtype="u1"),
        np.array([[[-1.0, 1.0], [np.nan, np.nan]], [[np.nan, np.nan], [np.nan, np.nan]]]),
        0.45,
    )


def test_world_cells_use_exact_pose_once_and_missing_is_not_free():
    replay = data()
    centers, sizes = world_cells(replay, 0, 2)
    np.testing.assert_allclose(centers, [[11.0, 22.0, 3.0]])
    np.testing.assert_allclose(sizes, [[0.45, 0.45, 2.0]])
    assert world_cells(replay, 0, 0)[0].shape == (0, 3)
    assert world_cells(replay, 1, 2)[0].shape == (0, 3)


def test_tgs_native_serialization_clears_missing_and_coverage_end(tmp_path):
    path = tmp_path / "scene.rrd"
    recording = rr.RecordingStream("nodedc_mission_core_recorded", recording_id="test")
    calls = []
    native_log = recording.log

    def capture(entity, value):
        calls.append((entity, value))
        native_log(entity, value)

    recording.log = capture
    recording.set_sinks(rr.FileSink(path, write_footer=True))
    log_tgs(recording, data(), 2.1)
    recording.flush()
    recording.disconnect()
    assert path.read_bytes().startswith(b"RRF2")
    assert sum(isinstance(value, rr.Boxes3D) for _, value in calls) == 1
    assert all(isinstance(value, rr.Clear) for _, value in calls[3:])
    assert calls[-1][0] == "/world/costmap"
    assert not any(entity.startswith("/world/points") for entity, _ in calls)


def test_coverage_cannot_end_before_last_tgs_sample():
    with pytest.raises(PortableReplayError):
        log_tgs(SimpleNamespace(), data(), 2.0)


def test_artifact_corruption_and_symlinks_are_rejected(tmp_path):
    path = tmp_path / "artifact"
    path.write_bytes(b"first")
    sha = hashlib.sha256(b"first").hexdigest()
    assert verified_file(path, sha, 5) == path
    path.write_bytes(b"other")
    with pytest.raises(PortableReplayError):
        verified_file(path, sha, 5)
    link = tmp_path / "link"
    link.symlink_to(path)
    with pytest.raises(PortableReplayError):
        verified_file(link, hashlib.sha256(b"other").hexdigest(), 5)


def test_camera_proxy_keeps_demux_time_base_and_reuses_its_own_sealed_cache(tmp_path, monkeypatch):
    import subprocess

    source = tmp_path / "init.mp4"
    source.write_bytes(b"synthetic")
    epoch = SimpleNamespace(
        init_path=source,
        init_sha256=hashlib.sha256(b"synthetic").hexdigest(),
        init_byte_length=9,
        segments=[],
    )
    service = PortableReplayService(
        view=None,
        data_dir=tmp_path,
        media=None,
        recording_source=lambda _: None,
        ffmpeg_path=Path("ffmpeg"),
    )
    service.cache.mkdir()
    staging = tmp_path / "staging"
    staging.mkdir()
    calls = []

    def run(command, **kwargs):
        calls.append(command)
        assert command[command.index("-enc_time_base") + 1] == "-1"
        assert command[command.index("-fps_mode") + 1] == "passthrough"
        assert "-r" not in command and "-frames:v" not in command
        Path(command[-1]).write_bytes(b"synthetic-proxy")
        return SimpleNamespace(returncode=0)

    monkeypatch.setattr(subprocess, "run", run)
    first = service._video(epoch, staging)
    assert first.parent == service.cache
    assert service._video(epoch, staging) == first
    assert len(calls) == 1


def test_replay_api_get_never_prepares_and_range_is_generation_fenced(tmp_path):
    path = tmp_path / "replay.rrd"
    path.write_bytes(b"RRF2test")
    sha = hashlib.sha256(path.read_bytes()).hexdigest()
    artifact = CanonicalLabReplayArtifact(path, 8, sha)
    prepared = []

    def prepare(*args):
        prepared.append(args)
        return artifact

    app = FastAPI()
    app.include_router(
        build_portable_replay_router(
            SimpleNamespace(prepare=prepare, cached=lambda *args: artifact if prepared else None)
        )
    )
    with TestClient(app) as client:
        url = f"/api/v1/observatory/portable-results/{RESULT}/replays/{BASE}/recording.rrd"
        assert client.get(url, params={"generation": sha}).status_code == 409
        assert not prepared
        response = client.head(url)
        assert response.status_code == 200 and response.headers["etag"] == f'"{sha}"'
        assert response.headers["x-rerun-format"] == "RRF2"
        response = client.get(url, params={"generation": sha}, headers={"Range": "bytes=0-3"})
        assert response.status_code == 206 and response.content == b"RRF2"
        assert client.get(url, params={"generation": "c" * 64}).status_code == 412
        assert len(prepared) == 1


def test_warm_reopen_does_not_prepare_and_detects_same_size_corruption(tmp_path, monkeypatch):
    import json

    service = PortableReplayService(
        view=SimpleNamespace(read=lambda _: {"artifact_manifest_id": "d" * 64}),
        data_dir=tmp_path,
        media=None,
        recording_source=lambda _: None,
        ffmpeg_path=Path("ffmpeg"),
    )
    key, _ = service._key(RESULT, BASE)
    service.cache.mkdir()
    path = service.cache / ("e" * 64 + ".replay.rrd")
    path.write_bytes(b"RRF2test")
    sha = hashlib.sha256(path.read_bytes()).hexdigest()
    (service.cache / f"{key}.json").write_text(
        json.dumps({"key": key, "file": path.name, "sha256": sha, "byte_length": 8})
    )
    assert service.prepare(RESULT, BASE).sha256 == sha
    monkeypatch.setattr(hashlib, "file_digest", lambda *_: pytest.fail("warm byte rehash"))
    assert service.prepare(RESULT, BASE).sha256 == sha
    monkeypatch.undo()
    path.write_bytes(b"RRF2fail")
    with pytest.raises(PortableReplayError):
        service.cached(RESULT, BASE)