"""Base classes for object-detection inference engines."""

from __future__ import annotations

import logging
from abc import ABC, abstractmethod
from enum import Enum
from typing import (
    TYPE_CHECKING,
    Any,
    ClassVar,
    Dict,
    List,
    Literal,
    Optional,
    Type,
    get_args,
    get_origin,
)

from PIL.Image import Image
from pydantic import BaseModel, ConfigDict, Field, field_validator
from pydantic_core import PydanticUndefined

if TYPE_CHECKING:
    from docling.datamodel.stage_model_specs import EngineModelConfig

_log = logging.getLogger(__name__)


class ObjectDetectionEngineType(str, Enum):
    """Supported inference engine types for object-detection models."""

    ONNXRUNTIME = "onnxruntime"
    TRANSFORMERS = "transformers"
    API_KSERVE_V2 = "api_kserve_v2"

    @classmethod
    def is_remote_variant(cls, engine_type: ObjectDetectionEngineType) -> bool:
        """Check if an engine type is a remote API variant."""
        return engine_type in {cls.API_KSERVE_V2}


class BaseObjectDetectionEngineOptions(BaseModel):
    """Base configuration shared across object-detection engines."""

    model_config = ConfigDict(arbitrary_types_allowed=True)

    engine_type: ObjectDetectionEngineType = Field(
        description="Type of inference engine to use",
    )

    score_threshold: float = Field(
        default=0.3,
        description="Minimum confidence score to keep a detection (0.0 to 1.0)",
    )

    _registry: ClassVar[
        dict[ObjectDetectionEngineType, Type[BaseObjectDetectionEngineOptions]]
    ] = {}

    @classmethod
    def __pydantic_init_subclass__(cls, **kwargs):
        super().__pydantic_init_subclass__(**kwargs)

        # Skip base class itself
        if cls is BaseObjectDetectionEngineOptions:
            return

        # only register concrete subclasses that fix engine_type via Literal
        field = cls.model_fields.get("engine_type")
        if not field:
            return

        engine_type = None

        # 1. Literal[...] annotation
        ann = field.annotation
        if get_origin(ann) is Literal:
            values = get_args(ann)
            if len(values) == 1:
                engine_type = values[0]

        # 2. Explicit default
        if engine_type is None and field.default is not PydanticUndefined:
            engine_type = field.default

        if engine_type is not None:
            BaseObjectDetectionEngineOptions._registry[engine_type] = cls


class ObjectDetectionEngineOptionsMixin(BaseModel):
    engine_options: BaseObjectDetectionEngineOptions = Field(
        description="Runtime configuration for the object-detection engine",
    )

    @field_validator("engine_options", mode="before")
    @classmethod
    def resolve_engine_options(cls, value):
        # already concrete
        if isinstance(value, BaseObjectDetectionEngineOptions):
            return value

        # dict / JSON case
        if isinstance(value, dict):
            engine_type = value.get("engine_type")
            model_cls = BaseObjectDetectionEngineOptions._registry.get(engine_type)
            if model_cls:
                return model_cls.model_validate(value)

        return value


class ObjectDetectionEngineInput(BaseModel):
    """Generic input accepted by every object-detection engine."""

    model_config = ConfigDict(arbitrary_types_allowed=True)

    image: Image = Field(description="PIL image to run inference on")
    metadata: Dict[str, Any] = Field(
        default_factory=dict,
        description="Optional metadata that is echoed back in the output",
    )


class ObjectDetectionEngineOutput(BaseModel):
    """Output returned by object-detection engines."""

    label_ids: List[int] = Field(
        default_factory=list,
        description="Predicted class indices",
    )
    scores: List[float] = Field(
        default_factory=list,
        description="Confidence scores for the predictions",
    )
    bboxes: List[List[float]] = Field(
        default_factory=list,
        description="Bounding boxes as [x_min, y_min, x_max, y_max] in pixels",
    )
    metadata: Dict[str, Any] = Field(
        default_factory=dict,
        description="Optional metadata echoed back from the input or engine",
    )


class BaseObjectDetectionEngine(ABC):
    """Abstract base-class for object-detection engines."""

    def __init__(
        self,
        options: BaseObjectDetectionEngineOptions,
        model_config: Optional[EngineModelConfig] = None,
    ) -> None:
        """Initialize the engine.

        Args:
            options: Engine-specific configuration options
            model_config: Model configuration (repo_id, revision, extra_config)
        """
        self.options = options
        self.model_config = model_config
        self._initialized = False

    @abstractmethod
    def initialize(self) -> None:
        """Initialize engine resources (load models, allocate buffers, etc.)."""

    @abstractmethod
    def predict_batch(
        self, input_batch: List[ObjectDetectionEngineInput]
    ) -> List[ObjectDetectionEngineOutput]:
        """Run inference on a batch of inputs."""

    @abstractmethod
    def get_label_mapping(self) -> Dict[int, str]:
        """Get the label mapping for this model.

        Returns:
            Dictionary mapping label IDs to label names
        """

    def predict(
        self, input_data: ObjectDetectionEngineInput
    ) -> ObjectDetectionEngineOutput:
        """Helper to run inference on a single input."""
        if not self._initialized:
            _log.debug("Initializing %s for single prediction", type(self).__name__)
            self.initialize()

        results = self.predict_batch([input_data])
        return results[0]

    def __call__(
        self,
        input_data: ObjectDetectionEngineInput | List[ObjectDetectionEngineInput],
    ) -> ObjectDetectionEngineOutput | List[ObjectDetectionEngineOutput]:
        if not self._initialized:
            _log.debug("Initializing %s for call", type(self).__name__)
            self.initialize()

        if isinstance(input_data, list):
            return self.predict_batch(input_data)
        return self.predict(input_data)
