"""Versioned fixed-class detector providers for raw-KB4 object proposals."""

from __future__ import annotations

import time
from collections import Counter
from collections.abc import Callable
from dataclasses import dataclass
from threading import Lock
from typing import Final

import numpy as np
from numpy.typing import NDArray

from .contracts import BoundingRegion2D, ObjectProposal2D
from .providers import SourcePacket
from .rf_detr_native_object_detector import (
    RF_DETR_NATIVE_CONFIG,
    RF_DETR_NATIVE_MODEL_ID,
    RF_DETR_NATIVE_MODEL_VERSION,
    NativeRfDetrConfig,
    NativeRfDetrInferenceBackend,
    postprocess_native_rf_detr,
    prepare_raw_kb4_rf_detr_native,
)
from .rf_detr_object_detector import (
    RF_DETR_CONFIG,
    RF_DETR_MODEL_ID,
    RF_DETR_MODEL_VERSION,
    RfDetrConfig,
    RfDetrDetection,
    RfDetrInferenceBackend,
    postprocess_rf_detr,
    preprocess_raw_kb4_rf_detr,
)
from .yolox_object_detector import (
    ALL_COCO_YOLOX_CONFIG,
    FROZEN_YOLOX_CONFIG,
    YOLOX_MODEL_ID,
    YOLOX_MODEL_VERSION,
    AllCocoYoloxConfig,
    FrozenYoloxConfig,
    ImageResizer,
    InferenceBackend,
    YoloxDetection,
    postprocess_yolox,
    preprocess_raw_kb4,
)

FROZEN_YOLOX_PROVIDER_ID: Final = "triton-yolox-s-raw-kb4/v1"
ALL_COCO_YOLOX_PROVIDER_ID: Final = "triton-yolox-s-raw-kb4-all-coco/v2"
FROZEN_YOLOX_MODEL_ID: Final = f"{YOLOX_MODEL_ID}:{YOLOX_MODEL_VERSION}"
FROZEN_YOLOX_PREPROCESS_ID: Final = "raw-kb4-valid-fov-letterbox/v1"
RF_DETR_SHADOW_PROVIDER_ID: Final = "triton-rf-detr-large-coco-risk-fp16-shadow/v0"
RF_DETR_SHADOW_MODEL_ID: Final = f"{RF_DETR_MODEL_ID}:{RF_DETR_MODEL_VERSION}"
RF_DETR_SHADOW_PREPROCESS_ID: Final = "raw-kb4-valid-fov-rgb-stretch-imagenet/v0"
RF_DETR_NATIVE_SHADOW_PROVIDER_ID: Final = (
    "triton-rf-detr-large-coco-native-kb4-risk-fp16-shadow/v0"
)
RF_DETR_NATIVE_SHADOW_MODEL_ID: Final = (
    f"{RF_DETR_NATIVE_MODEL_ID}:{RF_DETR_NATIVE_MODEL_VERSION}"
)
RF_DETR_NATIVE_SHADOW_PREPROCESS_ID: Final = (
    "raw-kb4-uint8-fused-mask-rgb-pad8-imagenet-trt/v0"
)


class DetectorProviderError(RuntimeError):
    """The detector input, frozen inference or proposal output is incompatible."""


@dataclass(frozen=True, slots=True)
class DetectorProviderSnapshot:
    input_frames: int
    completed_frames: int
    failed_frames: int
    zero_proposal_frames: int
    proposal_count: int
    rejected: tuple[tuple[str, int], ...]
    core_duration_ns: int


@dataclass(frozen=True, slots=True)
class DetectorFrameTiming:
    sequence: int
    preprocess_duration_ns: int
    inference_transport_duration_ns: int
    postprocess_duration_ns: int
    total_duration_ns: int

    def __post_init__(self) -> None:
        values = (
            self.sequence,
            self.preprocess_duration_ns,
            self.inference_transport_duration_ns,
            self.postprocess_duration_ns,
            self.total_duration_ns,
        )
        if any(value < 0 for value in values):
            raise DetectorProviderError("detector frame timing must be nonnegative")
        if (
            self.preprocess_duration_ns
            + self.inference_transport_duration_ns
            + self.postprocess_duration_ns
            != self.total_duration_ns
        ):
            raise DetectorProviderError("detector frame timing does not close")

    def to_dict(self) -> dict[str, int]:
        return {
            "sequence": self.sequence,
            "preprocess_duration_ns": self.preprocess_duration_ns,
            "inference_transport_duration_ns": self.inference_transport_duration_ns,
            "postprocess_duration_ns": self.postprocess_duration_ns,
            "total_duration_ns": self.total_duration_ns,
        }


DetectorTimingObserver = Callable[[DetectorFrameTiming], None]


@dataclass(frozen=True, slots=True)
class DetectorWarmupSnapshot:
    completed: bool
    inference_passes: int
    preprocess_duration_ns: int
    inference_transport_duration_ns: int
    postprocess_duration_ns: int
    total_duration_ns: int

    def __post_init__(self) -> None:
        durations = (
            self.preprocess_duration_ns,
            self.inference_transport_duration_ns,
            self.postprocess_duration_ns,
            self.total_duration_ns,
        )
        if (
            self.completed is not True
            or self.inference_passes != 1
            or any(value < 0 for value in durations)
            or sum(durations[:3]) != self.total_duration_ns
        ):
            raise DetectorProviderError("detector warmup snapshot is incompatible")


class FrozenYoloxDetectorProvider:
    """One image payload produces one frozen inference request and proposal tuple."""

    provider_id: str = FROZEN_YOLOX_PROVIDER_ID

    def __init__(
        self,
        *,
        mask: NDArray[np.bool_],
        backend: InferenceBackend,
        resizer: ImageResizer | None = None,
        config: FrozenYoloxConfig | AllCocoYoloxConfig = FROZEN_YOLOX_CONFIG,
        clock_ns: Callable[[], int] = time.perf_counter_ns,
    ) -> None:
        if mask.shape != (600, 800) or mask.dtype != np.bool_ or not np.any(mask):
            raise DetectorProviderError("frozen valid-FOV mask is incompatible")
        self.mask = np.asarray(mask, dtype=np.bool_)
        self.backend = backend
        self.resizer = resizer
        self.config = config
        self._clock_ns = clock_ns
        self._lock = Lock()
        self._input_frames = 0
        self._completed_frames = 0
        self._failed_frames = 0
        self._zero_proposal_frames = 0
        self._proposal_count = 0
        self._rejected: Counter[str] = Counter()
        self._core_duration_ns = 0

    def detect(self, packet: SourcePacket) -> tuple[ObjectProposal2D, ...]:
        payload = packet.image_payload
        with self._lock:
            self._input_frames += 1
        started_ns = int(self._clock_ns())
        try:
            if not isinstance(payload, np.ndarray):
                raise DetectorProviderError("detector requires a decoded BGR image payload")
            image = np.asarray(payload)
            if image.dtype != np.uint8:
                raise DetectorProviderError("decoded BGR image must be uint8")
            tensor = preprocess_raw_kb4(
                image,
                self.mask,
                config=self.config,
                resizer=self.resizer,
            )
            output = self.backend.infer(tensor)
            postprocessed = postprocess_yolox(output, self.mask, config=self.config)
            proposals = proposals_from_detections(
                packet,
                postprocessed.detections,
                provider_id=self.provider_id,
            )
        except Exception:
            with self._lock:
                self._failed_frames += 1
                self._core_duration_ns += max(0, int(self._clock_ns()) - started_ns)
            raise
        with self._lock:
            self._completed_frames += 1
            self._proposal_count += len(proposals)
            self._zero_proposal_frames += not proposals
            self._rejected.update(dict(postprocessed.rejected))
            self._core_duration_ns += max(0, int(self._clock_ns()) - started_ns)
        return proposals

    def snapshot(self) -> DetectorProviderSnapshot:
        with self._lock:
            return DetectorProviderSnapshot(
                input_frames=self._input_frames,
                completed_frames=self._completed_frames,
                failed_frames=self._failed_frames,
                zero_proposal_frames=self._zero_proposal_frames,
                proposal_count=self._proposal_count,
                rejected=tuple(sorted(self._rejected.items())),
                core_duration_ns=self._core_duration_ns,
            )


def proposals_from_detections(
    packet: SourcePacket,
    detections: tuple[YoloxDetection, ...],
    *,
    provider_id: str = FROZEN_YOLOX_PROVIDER_ID,
) -> tuple[ObjectProposal2D, ...]:
    envelope = packet.envelope
    return tuple(
        ObjectProposal2D(
            proposal_id=f"proposal-{envelope.sequence}-{index}",
            source_id=envelope.source_id,
            frame_id=envelope.frame_id,
            region=BoundingRegion2D(*detection.bbox_xyxy),
            objectness=detection.score,
            provider_id=provider_id,
            model_id=FROZEN_YOLOX_MODEL_ID,
            preprocess_id=FROZEN_YOLOX_PREPROCESS_ID,
            semantic_hint=detection.label,
            provider_tracklet=None,
        )
        for index, detection in enumerate(detections)
    )


class AllCocoYoloxDetectorProvider(FrozenYoloxDetectorProvider):
    """Emit every qualified COCO class without adding another inference pass."""

    provider_id: str = ALL_COCO_YOLOX_PROVIDER_ID

    def __init__(
        self,
        *,
        mask: NDArray[np.bool_],
        backend: InferenceBackend,
        resizer: ImageResizer | None = None,
        config: AllCocoYoloxConfig = ALL_COCO_YOLOX_CONFIG,
        clock_ns: Callable[[], int] = time.perf_counter_ns,
    ) -> None:
        super().__init__(
            mask=mask,
            backend=backend,
            resizer=resizer,
            config=config,
            clock_ns=clock_ns,
        )


class RfDetrShadowDetectorProvider:
    """Emit behavior-relevant fixed classes from one RF-DETR inference pass."""

    provider_id: str = RF_DETR_SHADOW_PROVIDER_ID

    def __init__(
        self,
        *,
        mask: NDArray[np.bool_],
        backend: RfDetrInferenceBackend,
        resizer: ImageResizer | None = None,
        config: RfDetrConfig = RF_DETR_CONFIG,
        clock_ns: Callable[[], int] = time.perf_counter_ns,
        timing_observer: DetectorTimingObserver | None = None,
    ) -> None:
        if mask.shape != (600, 800) or mask.dtype != np.bool_ or not np.any(mask):
            raise DetectorProviderError("RF-DETR valid-FOV mask is incompatible")
        self.mask = np.asarray(mask, dtype=np.bool_)
        self.backend = backend
        self.resizer = resizer
        self.config = config
        self._clock_ns = clock_ns
        self.timing_observer = timing_observer
        self._lock = Lock()
        self._input_frames = 0
        self._completed_frames = 0
        self._failed_frames = 0
        self._zero_proposal_frames = 0
        self._proposal_count = 0
        self._rejected: Counter[str] = Counter()
        self._core_duration_ns = 0
        self._warmup_started = False
        self._warmup_snapshot: DetectorWarmupSnapshot | None = None

    def warm_up(self) -> DetectorWarmupSnapshot:
        """Prime preprocessing, transport and postprocessing before source admission."""

        with self._lock:
            if self._warmup_snapshot is not None:
                return self._warmup_snapshot
            if self._warmup_started:
                raise DetectorProviderError("RF-DETR warmup is already in progress")
            self._warmup_started = True
        started_ns = int(self._clock_ns())
        try:
            image = np.zeros(
                (self.config.source_height, self.config.source_width, 3),
                dtype=np.uint8,
            )
            tensor = preprocess_raw_kb4_rf_detr(
                image,
                self.mask,
                config=self.config,
                resizer=self.resizer,
            )
            preprocessed_ns = int(self._clock_ns())
            output = self.backend.infer(tensor)
            inferred_ns = int(self._clock_ns())
            postprocess_rf_detr(output, self.mask, config=self.config)
            completed_ns = int(self._clock_ns())
        except Exception:
            with self._lock:
                self._warmup_started = False
            raise
        snapshot = DetectorWarmupSnapshot(
            completed=True,
            inference_passes=1,
            preprocess_duration_ns=max(0, preprocessed_ns - started_ns),
            inference_transport_duration_ns=max(0, inferred_ns - preprocessed_ns),
            postprocess_duration_ns=max(0, completed_ns - inferred_ns),
            total_duration_ns=max(0, completed_ns - started_ns),
        )
        with self._lock:
            self._warmup_snapshot = snapshot
        return snapshot

    def detect(self, packet: SourcePacket) -> tuple[ObjectProposal2D, ...]:
        payload = packet.image_payload
        with self._lock:
            self._input_frames += 1
        started_ns = int(self._clock_ns())
        try:
            if not isinstance(payload, np.ndarray):
                raise DetectorProviderError("RF-DETR requires a decoded BGR image payload")
            image = np.asarray(payload)
            if image.dtype != np.uint8:
                raise DetectorProviderError("decoded BGR image must be uint8")
            tensor = preprocess_raw_kb4_rf_detr(
                image,
                self.mask,
                config=self.config,
                resizer=self.resizer,
            )
            preprocessed_ns = (
                int(self._clock_ns()) if self.timing_observer is not None else started_ns
            )
            output = self.backend.infer(tensor)
            inferred_ns = (
                int(self._clock_ns()) if self.timing_observer is not None else preprocessed_ns
            )
            postprocessed = postprocess_rf_detr(output, self.mask, config=self.config)
            proposals = proposals_from_rf_detr_detections(packet, postprocessed.detections)
        except Exception:
            with self._lock:
                self._failed_frames += 1
                self._core_duration_ns += max(0, int(self._clock_ns()) - started_ns)
            raise
        completed_ns = int(self._clock_ns())
        with self._lock:
            self._completed_frames += 1
            self._proposal_count += len(proposals)
            self._zero_proposal_frames += not proposals
            self._rejected.update(dict(postprocessed.rejected))
            self._core_duration_ns += max(0, completed_ns - started_ns)
        if self.timing_observer is not None:
            self.timing_observer(
                DetectorFrameTiming(
                    sequence=packet.envelope.sequence,
                    preprocess_duration_ns=max(0, preprocessed_ns - started_ns),
                    inference_transport_duration_ns=max(0, inferred_ns - preprocessed_ns),
                    postprocess_duration_ns=max(0, completed_ns - inferred_ns),
                    total_duration_ns=max(0, completed_ns - started_ns),
                )
            )
        return proposals

    def snapshot(self) -> DetectorProviderSnapshot:
        with self._lock:
            return DetectorProviderSnapshot(
                input_frames=self._input_frames,
                completed_frames=self._completed_frames,
                failed_frames=self._failed_frames,
                zero_proposal_frames=self._zero_proposal_frames,
                proposal_count=self._proposal_count,
                rejected=tuple(sorted(self._rejected.items())),
                core_duration_ns=self._core_duration_ns,
            )


def proposals_from_rf_detr_detections(
    packet: SourcePacket,
    detections: tuple[RfDetrDetection, ...],
) -> tuple[ObjectProposal2D, ...]:
    envelope = packet.envelope
    return tuple(
        ObjectProposal2D(
            proposal_id=f"proposal-{envelope.sequence}-{index}",
            source_id=envelope.source_id,
            frame_id=envelope.frame_id,
            region=BoundingRegion2D(*detection.bbox_xyxy),
            objectness=detection.score,
            provider_id=RF_DETR_SHADOW_PROVIDER_ID,
            model_id=RF_DETR_SHADOW_MODEL_ID,
            preprocess_id=RF_DETR_SHADOW_PREPROCESS_ID,
            semantic_hint=detection.label,
            provider_tracklet=None,
        )
        for index, detection in enumerate(detections)
    )


class NativeRfDetrShadowDetectorProvider:
    """Emit risk classes from one exact-raster native RF-DETR inference pass."""

    provider_id: str = RF_DETR_NATIVE_SHADOW_PROVIDER_ID

    def __init__(
        self,
        *,
        mask: NDArray[np.bool_],
        backend: NativeRfDetrInferenceBackend,
        config: NativeRfDetrConfig = RF_DETR_NATIVE_CONFIG,
        clock_ns: Callable[[], int] = time.perf_counter_ns,
        timing_observer: DetectorTimingObserver | None = None,
    ) -> None:
        if mask.shape != (600, 800) or mask.dtype != np.bool_ or not np.any(mask):
            raise DetectorProviderError("native RF-DETR valid-FOV mask is incompatible")
        self.mask = np.asarray(mask, dtype=np.bool_)
        self.backend = backend
        self.config = config
        self._clock_ns = clock_ns
        self.timing_observer = timing_observer
        self._lock = Lock()
        self._input_frames = 0
        self._completed_frames = 0
        self._failed_frames = 0
        self._zero_proposal_frames = 0
        self._proposal_count = 0
        self._rejected: Counter[str] = Counter()
        self._core_duration_ns = 0
        self._warmup_started = False
        self._warmup_snapshot: DetectorWarmupSnapshot | None = None

    def warm_up(self) -> DetectorWarmupSnapshot:
        """Prime raw transport and postprocessing before source admission."""

        with self._lock:
            if self._warmup_snapshot is not None:
                return self._warmup_snapshot
            if self._warmup_started:
                raise DetectorProviderError("native RF-DETR warmup is already in progress")
            self._warmup_started = True
        started_ns = int(self._clock_ns())
        try:
            image = np.zeros(
                (self.config.source_height, self.config.source_width, 3),
                dtype=np.uint8,
            )
            tensor = prepare_raw_kb4_rf_detr_native(image, config=self.config)
            preprocessed_ns = int(self._clock_ns())
            output = self.backend.infer(tensor)
            inferred_ns = int(self._clock_ns())
            postprocess_native_rf_detr(output, self.mask, config=self.config)
            completed_ns = int(self._clock_ns())
        except Exception:
            with self._lock:
                self._warmup_started = False
            raise
        snapshot = DetectorWarmupSnapshot(
            completed=True,
            inference_passes=1,
            preprocess_duration_ns=max(0, preprocessed_ns - started_ns),
            inference_transport_duration_ns=max(0, inferred_ns - preprocessed_ns),
            postprocess_duration_ns=max(0, completed_ns - inferred_ns),
            total_duration_ns=max(0, completed_ns - started_ns),
        )
        with self._lock:
            self._warmup_snapshot = snapshot
        return snapshot

    def detect(self, packet: SourcePacket) -> tuple[ObjectProposal2D, ...]:
        payload = packet.image_payload
        with self._lock:
            self._input_frames += 1
        started_ns = int(self._clock_ns())
        try:
            if not isinstance(payload, np.ndarray):
                raise DetectorProviderError(
                    "native RF-DETR requires a decoded BGR image payload"
                )
            image = np.asarray(payload)
            if image.dtype != np.uint8:
                raise DetectorProviderError("decoded BGR image must be uint8")
            tensor = prepare_raw_kb4_rf_detr_native(image, config=self.config)
            preprocessed_ns = (
                int(self._clock_ns()) if self.timing_observer is not None else started_ns
            )
            output = self.backend.infer(tensor)
            inferred_ns = (
                int(self._clock_ns()) if self.timing_observer is not None else preprocessed_ns
            )
            postprocessed = postprocess_native_rf_detr(
                output,
                self.mask,
                config=self.config,
            )
            proposals = proposals_from_native_rf_detr_detections(
                packet,
                postprocessed.detections,
            )
        except Exception:
            with self._lock:
                self._failed_frames += 1
                self._core_duration_ns += max(0, int(self._clock_ns()) - started_ns)
            raise
        completed_ns = int(self._clock_ns())
        with self._lock:
            self._completed_frames += 1
            self._proposal_count += len(proposals)
            self._zero_proposal_frames += not proposals
            self._rejected.update(dict(postprocessed.rejected))
            self._core_duration_ns += max(0, completed_ns - started_ns)
        if self.timing_observer is not None:
            self.timing_observer(
                DetectorFrameTiming(
                    sequence=packet.envelope.sequence,
                    preprocess_duration_ns=max(0, preprocessed_ns - started_ns),
                    inference_transport_duration_ns=max(0, inferred_ns - preprocessed_ns),
                    postprocess_duration_ns=max(0, completed_ns - inferred_ns),
                    total_duration_ns=max(0, completed_ns - started_ns),
                )
            )
        return proposals

    def snapshot(self) -> DetectorProviderSnapshot:
        with self._lock:
            return DetectorProviderSnapshot(
                input_frames=self._input_frames,
                completed_frames=self._completed_frames,
                failed_frames=self._failed_frames,
                zero_proposal_frames=self._zero_proposal_frames,
                proposal_count=self._proposal_count,
                rejected=tuple(sorted(self._rejected.items())),
                core_duration_ns=self._core_duration_ns,
            )


def proposals_from_native_rf_detr_detections(
    packet: SourcePacket,
    detections: tuple[RfDetrDetection, ...],
) -> tuple[ObjectProposal2D, ...]:
    envelope = packet.envelope
    return tuple(
        ObjectProposal2D(
            proposal_id=f"proposal-{envelope.sequence}-{index}",
            source_id=envelope.source_id,
            frame_id=envelope.frame_id,
            region=BoundingRegion2D(*detection.bbox_xyxy),
            objectness=detection.score,
            provider_id=RF_DETR_NATIVE_SHADOW_PROVIDER_ID,
            model_id=RF_DETR_NATIVE_SHADOW_MODEL_ID,
            preprocess_id=RF_DETR_NATIVE_SHADOW_PREPROCESS_ID,
            semantic_hint=detection.label,
            provider_tracklet=None,
        )
        for index, detection in enumerate(detections)
    )


__all__ = [
    "ALL_COCO_YOLOX_PROVIDER_ID",
    "FROZEN_YOLOX_MODEL_ID",
    "FROZEN_YOLOX_PREPROCESS_ID",
    "FROZEN_YOLOX_PROVIDER_ID",
    "RF_DETR_SHADOW_MODEL_ID",
    "RF_DETR_SHADOW_PREPROCESS_ID",
    "RF_DETR_SHADOW_PROVIDER_ID",
    "RF_DETR_NATIVE_SHADOW_MODEL_ID",
    "RF_DETR_NATIVE_SHADOW_PREPROCESS_ID",
    "RF_DETR_NATIVE_SHADOW_PROVIDER_ID",
    "DetectorProviderError",
    "DetectorProviderSnapshot",
    "DetectorFrameTiming",
    "DetectorTimingObserver",
    "DetectorWarmupSnapshot",
    "AllCocoYoloxDetectorProvider",
    "FrozenYoloxDetectorProvider",
    "NativeRfDetrShadowDetectorProvider",
    "RfDetrShadowDetectorProvider",
    "proposals_from_detections",
    "proposals_from_native_rf_detr_detections",
    "proposals_from_rf_detr_detections",
]