"""Base classes for VLM inference engines."""

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 VlmEngineType(str, Enum):
    """Types of VLM inference engines available."""

    # Local/inline engines
    TRANSFORMERS = "transformers"
    MLX = "mlx"
    VLLM = "vllm"

    # API-based engines
    API = "api"
    API_OLLAMA = "api_ollama"
    API_LMSTUDIO = "api_lmstudio"
    API_OPENAI = "api_openai"

    # Auto-selection
    AUTO_INLINE = "auto_inline"

    @classmethod
    def is_api_variant(cls, engine_type: "VlmEngineType") -> bool:
        """Check if an engine type is an API variant."""
        return engine_type in {
            cls.API,
            cls.API_OLLAMA,
            cls.API_LMSTUDIO,
            cls.API_OPENAI,
        }

    @classmethod
    def is_inline_variant(cls, engine_type: "VlmEngineType") -> bool:
        """Check if an engine type is an inline/local variant."""
        return engine_type in {
            cls.TRANSFORMERS,
            cls.MLX,
            cls.VLLM,
        }


class BaseVlmEngineOptions(BaseModel):
    """Base configuration for VLM inference engines.

    Engine options are independent of model specifications and prompts.
    They only control how the inference is executed.
    """

    model_config = ConfigDict(arbitrary_types_allowed=True)

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

    # registry: engine_type → subclass
    _registry: ClassVar[Dict[VlmEngineType, Type["BaseVlmEngineOptions"]]] = {}

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

        # Skip base class itself
        if cls is BaseVlmEngineOptions:
            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:
            BaseVlmEngineOptions._registry[engine_type] = cls


class VlmEngineOptionsMixin(BaseModel):
    engine_options: BaseVlmEngineOptions = Field(
        description="Runtime configuration (transformers, mlx, api, etc.)"
    )

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

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

            # Minimal fallback: API variants share one options class
            if VlmEngineType.is_api_variant(engine_type):
                from docling.datamodel.vlm_engine_options import ApiVlmEngineOptions

                return ApiVlmEngineOptions.model_validate(value)

        return value


class VlmEngineInput(BaseModel):
    """Input to a VLM inference engine.

    This is the generic interface that all engines accept.
    """

    model_config = ConfigDict(arbitrary_types_allowed=True)

    image: Image = Field(description="PIL Image to process")
    prompt: str = Field(description="Text prompt for the model")
    temperature: float = Field(
        default=0.0, description="Sampling temperature for generation"
    )
    max_new_tokens: int = Field(
        default=4096, description="Maximum number of tokens to generate"
    )
    stop_strings: List[str] = Field(
        default_factory=list, description="Strings that trigger generation stopping"
    )
    extra_generation_config: Dict[str, Any] = Field(
        default_factory=dict, description="Additional generation configuration"
    )


class VlmEngineOutput(BaseModel):
    """Output from a VLM inference engine.

    This is the generic interface that all engines return.
    """

    text: str = Field(description="Generated text from the model")
    stop_reason: str | None = Field(
        default=None, description="Reason why generation stopped"
    )
    metadata: Dict[str, Any] = Field(
        default_factory=dict, description="Additional metadata from the engine"
    )


class BaseVlmEngine(ABC):
    """Abstract base class for VLM inference engines.

    An engine handles the low-level model inference with generic inputs
    (PIL images + text prompts) and returns text predictions.

    Engines are independent of:
    - Pipeline stages (DoclingDocument, Page objects)
    - Response formats (doctags, markdown, etc.)

    But they ARE aware of:
    - Model specifications (repo_id, revision, model_type via EngineModelConfig)

    These model specs are provided at construction time for eager initialization.
    """

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

        Args:
            options: Engine-specific configuration options
            model_config: Model configuration (repo_id, revision, extra_config)
                         If None, model must be specified in predict() calls
        """
        self.options = options
        self.model_config = model_config
        self._initialized: bool = False

    @abstractmethod
    def initialize(self) -> None:
        """Initialize the engine (load models, setup connections, etc.).

        This is called once before the first inference.
        Implementations should set self._initialized = True when done.
        """

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

        This is the primary method that all engines must implement.
        Single predictions are routed through this method.

        Args:
            input_batch: List of inputs to process

        Returns:
            List of outputs, one per input
        """

    def predict(self, input_data: VlmEngineInput) -> VlmEngineOutput:
        """Run inference on a single input.

        This is a convenience method that wraps the input in a list and calls
        predict_batch(). Engines should NOT override this method - all
        inference logic should be in predict_batch().

        Args:
            input_data: Generic input containing image, prompt, and config

        Returns:
            Generic output containing generated text and metadata
        """
        if not self._initialized:
            self.initialize()

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

    def __call__(
        self, input_data: VlmEngineInput | List[VlmEngineInput]
    ) -> VlmEngineOutput | List[VlmEngineOutput]:
        """Convenience method to run inference.

        Args:
            input_data: Single input or list of inputs

        Returns:
            Single output or list of outputs
        """
        if not self._initialized:
            self.initialize()

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

    def cleanup(self) -> None:
        """Clean up resources (optional).

        Called when the engine is no longer needed.
        Implementations can override to release resources.
        """
