from __future__ import annotations

import asyncio
import stat
import threading
import time
from pathlib import Path

import numpy as np
import pytest

from k1link.compute.live_perception import (
    LivePerceptionIngress,
    encode_live_perception_result,
)
from k1link.device_plugins.xgrids_k1.live_perception_shadow import (
    build_live_perception_result_receiver,
    build_live_perception_shadow_router,
    ensure_live_shadow_token,
)


def test_shadow_token_is_stable_and_private(
    tmp_path: Path,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.delenv("MISSIONCORE_DATA_DIR", raising=False)
    path, token = ensure_live_shadow_token(tmp_path)
    repeated_path, repeated_token = ensure_live_shadow_token(tmp_path)

    assert repeated_path == path
    assert repeated_token == token
    assert stat.S_IMODE(path.stat().st_mode) == 0o600
    assert stat.S_IMODE(path.parent.stat().st_mode) == 0o700


def test_shadow_router_exposes_only_the_exclusive_binary_stream() -> None:
    ingress = LivePerceptionIngress()
    router = build_live_perception_shadow_router(
        ingress,
        "test-plugin",
        bearer_token="x" * 43,
    )
    assert len(router.routes) == 1
    assert router.routes[0].path == ("/api/v1/device-plugins/test-plugin/live-perception-shadow")


def test_shadow_router_accepts_only_validated_diagnostic_results_back() -> None:
    ingress = LivePerceptionIngress()
    received: list[bytes] = []
    published = threading.Event()

    def receive_result(payload: bytes) -> bool:
        received.append(payload)
        published.set()
        return True

    router = build_live_perception_shadow_router(
        ingress,
        "test-plugin",
        bearer_token="x" * 43,
        result_receiver=receive_result,
    )
    ingress.begin_session("session-1")
    encoded = encode_live_perception_result(
        session_id="session-1",
        session_generation=1,
        frame_index=0,
        source_frame_index=0,
        session_seconds=0.0,
        captured_at_epoch_ns=1,
        image_jpeg=bytes.fromhex("ffd878ffd9"),
        segmentation_mask=np.zeros((600, 800), dtype=np.uint8),
        objects=[],
        delivery={"health": "healthy"},
    )

    class FakeWebSocket:
        def __init__(self) -> None:
            self.headers = {"authorization": f"Bearer {'x' * 43}"}
            self.messages = [
                {"type": "websocket.receive", "bytes": encoded},
                {"type": "websocket.disconnect"},
            ]
            self.sent: list[bytes] = []
            self.accepted = False

        async def accept(self) -> None:
            self.accepted = True

        async def receive(self) -> dict[str, object]:
            await asyncio.sleep(0)
            return self.messages.pop(0)

        async def send_bytes(self, payload: bytes) -> None:
            self.sent.append(payload)

        async def close(self, **_: object) -> None:
            return

    websocket = FakeWebSocket()
    asyncio.run(router.routes[0].endpoint(websocket))  # type: ignore[attr-defined]

    assert websocket.accepted is True
    assert published.is_set()
    assert received == [encoded]


def test_shadow_result_receiver_rejects_previous_acquisition_generation() -> None:
    ingress = LivePerceptionIngress()
    published: list[tuple[str, int]] = []
    receiver = build_live_perception_result_receiver(
        ingress,
        lambda frame: not published.append((frame.session_id, frame.session_generation)),
    )
    ingress.begin_session("session-1")
    encoded = encode_live_perception_result(
        session_id="session-1",
        session_generation=1,
        frame_index=0,
        source_frame_index=0,
        session_seconds=0.0,
        captured_at_epoch_ns=1,
        image_jpeg=bytes.fromhex("ffd878ffd9"),
        segmentation_mask=None,
        objects=[],
        delivery={"health": "healthy"},
    )
    ingress.end_session("session-1")
    ingress.begin_session("session-2")

    assert receiver(encoded) is False
    assert published == []
    snapshot = ingress.snapshot()
    assert snapshot["results_accepted"] == 0
    assert snapshot["results_rejected_stale"] == 1
    assert snapshot["results_rejected_receiver"] == 0


def test_shadow_router_closes_a_worker_that_publishes_for_previous_session() -> None:
    ingress = LivePerceptionIngress()
    published: list[int] = []
    receiver = build_live_perception_result_receiver(
        ingress,
        lambda frame: not published.append(frame.frame_index),
    )
    router = build_live_perception_shadow_router(
        ingress,
        "test-plugin",
        bearer_token="x" * 43,
        result_receiver=receiver,
    )
    ingress.begin_session("session-1")
    stale = encode_live_perception_result(
        session_id="session-1",
        session_generation=1,
        frame_index=7,
        source_frame_index=7,
        session_seconds=0.0,
        captured_at_epoch_ns=1,
        image_jpeg=bytes.fromhex("ffd878ffd9"),
        segmentation_mask=None,
        objects=[],
        delivery={"health": "healthy"},
    )
    ingress.end_session("session-1")
    ingress.begin_session("session-2")

    class StaleResultWebSocket:
        def __init__(self) -> None:
            self.headers = {"authorization": f"Bearer {'x' * 43}"}
            self.closed: list[dict[str, object]] = []

        async def accept(self) -> None:
            return

        async def receive(self) -> dict[str, object]:
            await asyncio.sleep(0)
            return {"type": "websocket.receive", "bytes": stale}

        async def send_bytes(self, _payload: bytes) -> None:
            return

        async def close(self, **values: object) -> None:
            self.closed.append(values)

    websocket = StaleResultWebSocket()
    asyncio.run(router.routes[0].endpoint(websocket))  # type: ignore[attr-defined]

    assert published == []
    assert any(
        item.get("code") == 1008 and item.get("reason") == "Shadow result session is stale"
        for item in websocket.closed
    )
    assert ingress.snapshot()["results_rejected_stale"] == 1


def test_shadow_router_does_not_discard_ingress_while_receiving_results() -> None:
    class SlowIngress(LivePerceptionIngress):
        def take_next(
            self,
            consumer_id: str,
            *,
            timeout: float | None = None,
        ):
            time.sleep(0.01)
            return super().take_next(consumer_id, timeout=timeout)

    ingress = SlowIngress()
    received: list[bytes] = []
    router = build_live_perception_shadow_router(
        ingress,
        "test-plugin",
        bearer_token="x" * 43,
        result_receiver=lambda payload: not received.append(payload),
    )
    ingress.begin_session("session-1")
    for modality, sequence in (
        ("camera-init", 0),
        ("camera-frame", 1),
        ("lidar", 1),
        ("pose", 1),
    ):
        assert ingress.publish(
            modality=modality,
            source_id="sensor.camera.right",
            source_sequence=sequence,
            captured_at_epoch_ns=sequence,
            received_monotonic_ns=sequence,
            payload=modality.encode(),
        )
    ingress.end_session("session-1")
    expected_events = sum(int(queue["depth"]) for queue in ingress.snapshot()["queues"].values())

    class DuplexFakeWebSocket:
        def __init__(self) -> None:
            self.headers = {"authorization": f"Bearer {'x' * 43}"}
            self.results_remaining = 8
            self.sent: list[bytes] = []

        async def accept(self) -> None:
            return

        async def receive(self) -> dict[str, object]:
            if self.results_remaining:
                self.results_remaining -= 1
                await asyncio.sleep(0)
                return {"type": "websocket.receive", "bytes": b"result"}
            deadline = asyncio.get_running_loop().time() + 1
            while len(self.sent) < expected_events:
                if asyncio.get_running_loop().time() >= deadline:
                    break
                await asyncio.sleep(0.005)
            return {"type": "websocket.disconnect"}

        async def send_bytes(self, payload: bytes) -> None:
            self.sent.append(payload)

        async def close(self, **_: object) -> None:
            return

    websocket = DuplexFakeWebSocket()
    asyncio.run(router.routes[0].endpoint(websocket))  # type: ignore[attr-defined]

    assert len(received) == 8
    assert len(websocket.sent) == expected_events
    snapshot = ingress.snapshot()
    assert sum(int(queue["consumed"]) for queue in snapshot["queues"].values()) == expected_events