from __future__ import annotations

import asyncio
import logging
import threading
from collections.abc import Awaitable, Callable, Mapping, Sequence
from concurrent.futures import ThreadPoolExecutor
from typing import Any, TypeVar

import pytest

from cmvr_edge_ai.contracts import BoundingBox, Detection, DetectionResult, ImageFrame
from cmvr_edge_ai.core.component import ComponentContext, Emission
from cmvr_edge_ai.core.envelope import Envelope
from cmvr_edge_ai.detection.base import DetectionModel
from cmvr_edge_ai.detection.operator import DetectionOperator
from cmvr_edge_ai.detection.registry import (
    DetectionModelRegistry,
    DetectionModelSpec,
    UnknownDetectionModelError,
)


ResultT = TypeVar("ResultT")


class RecordingDetectionModel(DetectionModel):
    def __init__(self, detections: Sequence[Detection] = ()) -> None:
        self.detections = tuple(detections)
        self.load_calls = 0
        self.predict_calls = 0
        self.close_calls = 0
        self.predict_arguments: list[tuple[ImageFrame, tuple[str, ...], float]] = []
        self.call_threads: dict[str, list[int]] = {
            "load": [],
            "predict": [],
            "close": [],
        }

    def load(self) -> None:
        self.load_calls += 1
        self.call_threads["load"].append(threading.get_ident())

    def predict(
        self,
        frame: ImageFrame,
        labels: tuple[str, ...],
        confidence: float,
    ) -> Sequence[Detection]:
        self.predict_calls += 1
        self.predict_arguments.append((frame, labels, confidence))
        self.call_threads["predict"].append(threading.get_ident())
        return self.detections

    def close(self) -> None:
        self.close_calls += 1
        self.call_threads["close"].append(threading.get_ident())


def _detection(label: str, confidence: float) -> Detection:
    return Detection(
        label=label,
        confidence=confidence,
        box=BoundingBox(x_min=1.0, y_min=2.0, x_max=11.0, y_max=22.0),
    )


def _frame(*, codec: str = "none") -> ImageFrame:
    return ImageFrame(
        data=bytes(range(12)),
        width=2,
        height=2,
        pixel_format="BGR8",
        codec=codec,
        is_key_frame=codec != "none",
    )


def _envelope(frame: ImageFrame | None = None) -> Envelope[ImageFrame]:
    return Envelope(
        payload=_frame() if frame is None else frame,
        schema_name="ImageFrame",
        schema_version=1,
        source_id="camera-right",
        sequence=42,
        captured_at_ns=1_234_000_000,
        received_at_ns=1_235_000_000,
        deadline_ns=1_500_000_000,
        trace_id="trace-detection-42",
        session_id="camera-session",
        attributes={"site": "park-a"},
    )


def _registry(
    model: RecordingDetectionModel,
    *,
    labels: tuple[str, ...] = ("Worker", "No-Helmet", "No-Vest"),
    factory_options: list[dict[str, Any]] | None = None,
) -> DetectionModelRegistry:
    def factory(options: Mapping[str, Any]) -> DetectionModel:
        if factory_options is not None:
            factory_options.append(dict(options))
        return model

    registry = DetectionModelRegistry()
    registry.register(
        DetectionModelSpec(
            model_id="construction-ppe@1",
            name="Construction PPE",
            supported_labels=labels,
            factory=factory,
            backend="fake",
        )
    )
    return registry


def _context(executor: ThreadPoolExecutor) -> ComponentContext:
    return ComponentContext(
        pipeline_id="detection-test",
        node_id="detector",
        shutdown_event=asyncio.Event(),
        metadata={"thread_executor": executor},
    )


async def _keep_restricted_event_loop_responsive() -> None:
    """Provide wakeups where sandboxed self-pipe notifications are unavailable."""

    while True:
        await asyncio.sleep(0.01)


def _run_with_thread_executor(
    exercise: Callable[[ThreadPoolExecutor], Awaitable[ResultT]],
) -> ResultT:
    async def runner() -> ResultT:
        loop = asyncio.get_running_loop()
        executor = ThreadPoolExecutor(max_workers=1)
        loop.set_default_executor(executor)
        ticker = asyncio.create_task(_keep_restricted_event_loop_responsive())
        try:
            return await exercise(executor)
        finally:
            ticker.cancel()
            await asyncio.gather(ticker, return_exceptions=True)
            executor.shutdown(wait=True, cancel_futures=True)
            # asyncio.run() otherwise tries to shut down the same executor again.
            loop._default_executor = None  # type: ignore[attr-defined]

    return asyncio.run(runner())


def test_operator_loads_once_predicts_and_closes_idempotently() -> None:
    model = RecordingDetectionModel([_detection("Worker", 0.9)])
    options_seen: list[dict[str, Any]] = []
    operator = DetectionOperator(
        "detector",
        {
            "model": "construction-ppe@1",
            "model_options": {"weights": "/models/ppe.pt", "device": "cpu"},
        },
        model_registry=_registry(model, factory_options=options_seen),
    )

    async def scenario(executor: ThreadPoolExecutor) -> None:
        await operator.setup(_context(executor))
        result = await operator.process(_envelope())
        await operator.stop()
        await operator.stop()

        assert isinstance(result, Emission)
        assert model.load_calls == 1
        assert model.predict_calls == 1
        assert model.close_calls == 1
        assert options_seen == [{"weights": "/models/ppe.pt", "device": "cpu"}]

    _run_with_thread_executor(scenario)


def test_operator_filters_selected_labels_and_per_label_confidence() -> None:
    selected_worker = _detection("Worker", 0.50)
    selected_violation = _detection("No-Helmet", 0.80)
    model = RecordingDetectionModel(
        [
            _detection("Worker", 0.49),
            selected_worker,
            _detection("No-Helmet", 0.79),
            selected_violation,
            _detection("No-Vest", 0.99),
        ]
    )
    operator = DetectionOperator(
        "detector",
        {
            "model": "construction-ppe@1",
            "detect_labels": ["Worker", "No-Helmet"],
            "confidence": 0.50,
            "label_confidence": {"No-Helmet": 0.80},
        },
        model_registry=_registry(model),
    )

    async def scenario(executor: ThreadPoolExecutor) -> None:
        await operator.setup(_context(executor))
        emission = await operator.process(_envelope())
        await operator.stop()

        assert isinstance(emission, Emission)
        payload = emission.envelope.payload
        assert isinstance(payload, DetectionResult)
        assert payload.detections == (selected_worker, selected_violation)
        assert operator.selected_labels == ("Worker", "No-Helmet")
        assert model.predict_arguments[0][1:] == (("Worker", "No-Helmet"), 0.50)

    _run_with_thread_executor(scenario)


def test_operator_rejects_unknown_model_during_construction() -> None:
    with pytest.raises(UnknownDetectionModelError, match="missing@1"):
        DetectionOperator(
            "detector",
            {"model": "missing@1"},
            model_registry=DetectionModelRegistry(),
        )


@pytest.mark.parametrize(
    "params",
    [
        {
            "model": "construction-ppe@1",
            "detect_labels": ["Unknown-Label"],
        },
        {
            "model": "construction-ppe@1",
            "label_confidence": {"Unknown-Label": 0.9},
        },
    ],
)
def test_operator_rejects_unknown_labels_during_construction(
    params: dict[str, Any],
) -> None:
    with pytest.raises(ValueError, match="Unknown-Label"):
        DetectionOperator(
            "detector",
            params,
            model_registry=_registry(RecordingDetectionModel()),
        )


def test_operator_rejects_encoded_frame_without_calling_model() -> None:
    model = RecordingDetectionModel()
    operator = DetectionOperator(
        "detector",
        {"model": "construction-ppe@1"},
        model_registry=_registry(model),
    )

    async def scenario(executor: ThreadPoolExecutor) -> None:
        await operator.setup(_context(executor))
        with pytest.raises(ValueError, match="decoded ImageFrame"):
            await operator.process(_envelope(_frame(codec="h264")))
        await operator.stop()

        assert model.predict_calls == 0

    _run_with_thread_executor(scenario)


def test_operator_emits_result_metadata_and_preserves_envelope_correlation() -> None:
    source = _envelope()
    expected_detection = _detection("No-Helmet", 0.93)
    model = RecordingDetectionModel([expected_detection])
    operator = DetectionOperator(
        "detector",
        {"model": "construction-ppe@1"},
        model_registry=_registry(model),
    )

    async def scenario(executor: ThreadPoolExecutor) -> None:
        await operator.setup(_context(executor))
        emission = await operator.process(source)
        await operator.stop()

        assert isinstance(emission, Emission)
        assert emission.port == "detections"
        output = emission.envelope
        assert output.schema == "DetectionResult/v1"
        assert isinstance(output.payload, DetectionResult)
        assert output.payload.detections == (expected_detection,)
        assert output.payload.model_id == "construction-ppe@1"
        assert output.payload.model_name == "Construction PPE"
        assert output.payload.inference_ms >= 0.0
        assert output.payload.source_frame is None
        assert operator.attach_frame is False
        assert output.source_id == source.source_id
        assert output.sequence == source.sequence
        assert output.captured_at_ns == source.captured_at_ns
        assert output.received_at_ns == source.received_at_ns
        assert output.deadline_ns == source.deadline_ns
        assert output.trace_id == source.trace_id
        assert output.session_id == source.session_id
        assert output.attributes == {
            "site": "park-a",
            "detection_model_id": "construction-ppe@1",
            "detection_model_name": "Construction PPE",
        }

    _run_with_thread_executor(scenario)


def test_operator_attach_frame_keeps_the_decoded_frame_by_reference() -> None:
    source = _envelope()
    model = RecordingDetectionModel([_detection("No-Helmet", 0.93)])
    operator = DetectionOperator(
        "detector",
        {"model": "construction-ppe@1", "attach_frame": True},
        model_registry=_registry(model),
    )

    async def scenario(executor: ThreadPoolExecutor) -> None:
        await operator.setup(_context(executor))
        emission = await operator.process(source)
        await operator.stop()

        assert isinstance(emission, Emission)
        assert isinstance(emission.envelope.payload, DetectionResult)
        assert emission.envelope.payload.source_frame is source.payload
        assert operator.attach_frame is True

    _run_with_thread_executor(scenario)


def test_blocking_model_methods_run_on_the_supplied_thread_pool() -> None:
    event_loop_thread = threading.get_ident()
    model = RecordingDetectionModel()
    operator = DetectionOperator(
        "detector",
        {"model": "construction-ppe@1"},
        model_registry=_registry(model),
    )

    async def scenario(executor: ThreadPoolExecutor) -> None:
        await operator.setup(_context(executor))
        await operator.process(_envelope())
        await operator.stop()

        assert all(len(thread_ids) == 1 for thread_ids in model.call_threads.values())
        worker_threads = {
            thread_id
            for thread_ids in model.call_threads.values()
            for thread_id in thread_ids
        }
        assert len(worker_threads) == 1
        assert event_loop_thread not in worker_threads

    _run_with_thread_executor(scenario)


def test_max_fps_skips_frames_without_invoking_model_again() -> None:
    model = RecordingDetectionModel()
    operator = DetectionOperator(
        "detector",
        {"model": "construction-ppe@1", "max_fps": 0.001},
        model_registry=_registry(model),
    )

    async def scenario(executor: ThreadPoolExecutor) -> None:
        await operator.setup(_context(executor))
        first = await operator.process(_envelope())
        second = await operator.process(_envelope())
        await operator.stop()

        assert isinstance(first, Emission)
        assert second is None
        assert model.predict_calls == 1

    _run_with_thread_executor(scenario)


def test_operator_logs_model_load_and_configured_inference_heartbeat(
    caplog: pytest.LogCaptureFixture,
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    timestamps = iter((0, 1_000_000_000, 6_000_000_000))
    monkeypatch.setattr(
        "cmvr_edge_ai.detection.operator.monotonic_ns",
        lambda: next(timestamps),
    )
    model = RecordingDetectionModel([_detection("No-Helmet", 0.93)])
    operator = DetectionOperator(
        "detector",
        {
            "model": "construction-ppe@1",
            "detect_labels": ["No-Helmet"],
            "inference_log_interval_s": 5,
        },
        model_registry=_registry(model),
    )

    async def scenario(executor: ThreadPoolExecutor) -> None:
        await operator.setup(_context(executor))
        await operator.process(_envelope())
        await operator.process(_envelope())
        await operator.process(_envelope())
        await operator.stop()

    with caplog.at_level(logging.INFO, logger="cmvr_edge_ai.detection.operator"):
        _run_with_thread_executor(scenario)

    messages = [record.getMessage() for record in caplog.records]
    assert any(
        "detection model loaded node=detector model=construction-ppe@1" in message
        for message in messages
    )
    inference_messages = [
        message for message in messages if message.startswith("detection inference ")
    ]
    assert len(inference_messages) == 2
    assert "total_frames=1" in inference_messages[0]
    assert "window_frames=1" in inference_messages[0]
    assert "window_detections=1" in inference_messages[0]
    assert "hit_labels=No-Helmet:1" in inference_messages[0]
    assert "total_frames=3" in inference_messages[1]
    assert "window_frames=2" in inference_messages[1]
    assert "window_detections=2" in inference_messages[1]
    assert "hit_labels=No-Helmet:2" in inference_messages[1]


@pytest.mark.parametrize("value", [0, -1, 3600.1, True])
def test_operator_rejects_invalid_inference_log_interval(value: Any) -> None:
    with pytest.raises(ValueError, match="inference_log_interval_s"):
        DetectionOperator(
            "detector",
            {
                "model": "construction-ppe@1",
                "inference_log_interval_s": value,
            },
            model_registry=_registry(RecordingDetectionModel()),
        )
