from __future__ import annotations

import asyncio
from collections.abc import Mapping, Sequence
from concurrent.futures import ThreadPoolExecutor
from typing import Any

from cmvr_edge_ai.application import (
    create_default_model_registry,
    create_default_registry,
)
from cmvr_edge_ai.compiler import compile_pipeline
from cmvr_edge_ai.config import load_config_data
from cmvr_edge_ai.contracts import (
    BoundingBox,
    Detection,
    DetectionAlert,
    DetectionResult,
    ImageFrame,
)
from cmvr_edge_ai.core import Envelope, Sink
from cmvr_edge_ai.detection import (
    DetectionModel,
    DetectionModelRegistry,
    DetectionModelSpec,
)
from cmvr_edge_ai.plugins import PluginKind, PluginSpec


SIX_CLASS_MODEL_ID = "ppe-6classes-yolov8n@1"
SIX_CLASS_LABELS = (
    "Gloves",
    "Vest",
    "goggles",
    "helmet",
    "mask",
    "safety_shoe",
)


class _FixedDetectionModel(DetectionModel):
    def __init__(self, label: str) -> None:
        self.label = label
        self.load_calls = 0
        self.close_calls = 0
        self.frames: list[ImageFrame] = []
        self.selected_labels: list[tuple[str, ...]] = []

    def load(self) -> None:
        self.load_calls += 1

    def predict(
        self,
        frame: ImageFrame,
        labels: tuple[str, ...],
        confidence: float,
    ) -> Sequence[Detection]:
        del confidence
        self.frames.append(frame)
        self.selected_labels.append(labels)
        return (
            Detection(
                label=self.label,
                confidence=0.95,
                box=BoundingBox(1.0, 2.0, 20.0, 30.0),
            ),
        )

    def close(self) -> None:
        self.close_calls += 1


class _CollectSink(Sink):
    def __init__(self) -> None:
        self.envelopes: list[Envelope[Any]] = []
        self.input_ports: list[str] = []

    async def consume(
        self,
        envelope: Envelope[Any],
        input_port: str = "input",
    ) -> None:
        self.envelopes.append(envelope)
        self.input_ports.append(input_port)


async def _keep_event_loop_responsive() -> None:
    while True:
        await asyncio.sleep(0.01)


def test_default_registry_contains_six_class_ppe_model_contract() -> None:
    registry = create_default_model_registry(discover_entry_points=False)

    spec = registry.resolve(SIX_CLASS_MODEL_ID)

    assert spec.model_id == SIX_CLASS_MODEL_ID
    assert spec.supported_labels == SIX_CLASS_LABELS


def test_decoder_fans_out_to_alert_and_raw_detection_branches() -> None:
    ppe_model = _FixedDetectionModel("No-Helmet")
    six_class_model = _FixedDetectionModel("helmet")
    models = DetectionModelRegistry()
    models.register(
        DetectionModelSpec(
            model_id="construction-ppe-yolov8@1",
            name="Test Construction PPE",
            supported_labels=("No-Helmet",),
            factory=lambda options: ppe_model,
            backend="fake",
        )
    )
    models.register(
        DetectionModelSpec(
            model_id=SIX_CLASS_MODEL_ID,
            name="Test Six-Class PPE",
            supported_labels=SIX_CLASS_LABELS,
            factory=lambda options: six_class_model,
            backend="fake",
        )
    )
    registry = create_default_registry(
        discover_entry_points=False,
        model_registry=models,
    )
    alert_sink = _CollectSink()
    raw_sink = _CollectSink()
    registry.register(
        PluginSpec(
            plugin_id="test.dual_alert_sink@1",
            kind=PluginKind.SINK,
            factory=lambda node_id, params: alert_sink,
            inputs={"alerts": "DetectionAlert/v1"},
        )
    )
    registry.register(
        PluginSpec(
            plugin_id="test.dual_raw_detection_sink@1",
            kind=PluginKind.SINK,
            factory=lambda node_id, params: raw_sink,
            inputs={"detections": "DetectionResult/v1"},
        )
    )
    width, height = 4, 3
    frames = (
        ImageFrame(bytes((0, 0, 255)) * (width * height), width, height, "BGR8"),
        ImageFrame(bytes((0, 255, 0)) * (width * height), width, height, "BGR8"),
    )
    config = load_config_data(
        {
            "api_version": "cmvr.edge.ai/v1",
            "pipelines": {
                "dual_detection": {
                    "nodes": {
                        "source": {
                            "uses": "core.sequence_source@1",
                            "with": {
                                "items": list(frames),
                                "schema_name": "ImageFrame",
                            },
                        },
                        "decoder": {"uses": "media.video_decoder.pyav@1"},
                        "ppe_detector": {
                            "uses": "detection.model@1",
                            "with": {
                                "model": "construction-ppe-yolov8@1",
                                "detect_labels": ["No-Helmet"],
                            },
                        },
                        "ppe_gate": {
                            "uses": "detection.repeat_gate@1",
                            "with": {
                                "time_source": "received",
                                "rules": [
                                    {
                                        "id": "no-helmet",
                                        "labels": ["No-Helmet"],
                                        "min_hits": 2,
                                        "window_ms": 60_000,
                                        "cooldown_ms": 60_000,
                                        "scope": "source",
                                    }
                                ],
                            },
                        },
                        "alert_sink": {"uses": "test.dual_alert_sink@1"},
                        "six_class_detector": {
                            "uses": "detection.model@1",
                            "with": {
                                "model": SIX_CLASS_MODEL_ID,
                                "detect_labels": list(SIX_CLASS_LABELS),
                                "attach_frame": False,
                            },
                        },
                        "raw_sink": {
                            "uses": "test.dual_raw_detection_sink@1"
                        },
                    },
                    "edges": [
                        {
                            "from": "source.output",
                            "to": "decoder.frames",
                            "qos": {
                                "profile": "video_contiguous",
                                "capacity": 4,
                                "overflow": "block",
                            },
                        },
                        {
                            "from": "decoder.frames",
                            "to": "ppe_detector.frames",
                            "qos": {
                                "profile": "telemetry",
                                "capacity": 4,
                                "overflow": "block",
                            },
                        },
                        {
                            "from": "decoder.frames",
                            "to": "six_class_detector.frames",
                            "qos": {
                                "profile": "telemetry",
                                "capacity": 4,
                                "overflow": "block",
                            },
                        },
                        {
                            "from": "ppe_detector.detections",
                            "to": "ppe_gate.detections",
                            "qos": {
                                "profile": "telemetry",
                                "capacity": 4,
                                "overflow": "block",
                            },
                        },
                        {
                            "from": "ppe_gate.alerts",
                            "to": "alert_sink.alerts",
                            "qos": {
                                "profile": "telemetry",
                                "capacity": 4,
                                "overflow": "block",
                            },
                        },
                        {
                            "from": "six_class_detector.detections",
                            "to": "raw_sink.detections",
                            "qos": {
                                "profile": "telemetry",
                                "capacity": 4,
                                "overflow": "block",
                            },
                        },
                    ],
                }
            },
        }
    )

    async def scenario():  # type: ignore[no-untyped-def]
        loop = asyncio.get_running_loop()
        executor = ThreadPoolExecutor(max_workers=2)
        loop.set_default_executor(executor)
        ticker = asyncio.create_task(_keep_event_loop_responsive())
        try:
            compiled = compile_pipeline(
                config,
                "dual_detection",
                registry,
                metadata={"thread_executor": executor},
            )
            await compiled.runtime.run()
            return compiled
        finally:
            ticker.cancel()
            await asyncio.gather(ticker, return_exceptions=True)
            executor.shutdown(wait=True, cancel_futures=True)
            loop._default_executor = None  # type: ignore[attr-defined]

    compiled = asyncio.run(scenario())

    assert compiled.plugin_specs["ppe_gate"].outputs == {
        "alerts": "DetectionAlert/v1"
    }
    assert compiled.plugin_specs["six_class_detector"].outputs == {
        "detections": "DetectionResult/v1"
    }

    assert ppe_model.load_calls == ppe_model.close_calls == 1
    assert six_class_model.load_calls == six_class_model.close_calls == 1
    assert len(ppe_model.frames) == len(six_class_model.frames) == len(frames)
    assert all(
        ppe_frame is six_class_frame
        for ppe_frame, six_class_frame in zip(
            ppe_model.frames,
            six_class_model.frames,
        )
    )
    assert ppe_model.selected_labels == [("No-Helmet",)] * len(frames)
    assert six_class_model.selected_labels == [SIX_CLASS_LABELS] * len(frames)

    assert len(alert_sink.envelopes) == 1
    alert_envelope = alert_sink.envelopes[0]
    assert alert_sink.input_ports == ["alerts"]
    assert alert_envelope.schema == "DetectionAlert/v1"
    assert alert_envelope.source_id == "source"
    assert alert_envelope.sequence == 1
    assert isinstance(alert_envelope.payload, DetectionAlert)
    assert alert_envelope.payload.model_id == "construction-ppe-yolov8@1"
    assert alert_envelope.payload.rule_id == "no-helmet"

    assert raw_sink.input_ports == ["detections", "detections"]
    assert [envelope.schema for envelope in raw_sink.envelopes] == [
        "DetectionResult/v1",
        "DetectionResult/v1",
    ]
    assert [envelope.source_id for envelope in raw_sink.envelopes] == [
        "source",
        "source",
    ]
    assert [envelope.sequence for envelope in raw_sink.envelopes] == [0, 1]
    assert all(
        isinstance(envelope.payload, DetectionResult)
        for envelope in raw_sink.envelopes
    )
    assert all(
        envelope.payload.model_id == SIX_CLASS_MODEL_ID
        for envelope in raw_sink.envelopes
    )
    assert all(
        envelope.payload.source_frame is None
        for envelope in raw_sink.envelopes
    )
    assert alert_envelope.trace_id == raw_sink.envelopes[1].trace_id

    stats = compiled.runtime.edge_stats()
    ppe_fanout = next(
        value
        for name, value in stats.items()
        if name.endswith("decoder.frames->ppe_detector.frames")
    )
    six_class_fanout = next(
        value
        for name, value in stats.items()
        if name.endswith("decoder.frames->six_class_detector.frames")
    )
    assert ppe_fanout.enqueued == len(frames)
    assert six_class_fanout.enqueued == len(frames)
