#!/usr/bin/env python

"""████
 ██    ██    Datature
   ██  ██    Powering Breakthrough AI
     ██

@File    :   qwenvl.py
@Author  :   Wei Loon Cheng
@Contact :   developers@datature.io
@License :   Apache License 2.0
@Desc    :   Datature Vi SDK Qwen VL predictor module.
"""

from vi.inference.errors import wrap_configuration_error
from vi.inference.messages import (
    create_system_message,
    create_user_message_with_image,
    create_user_message_with_image_url,
    create_user_message_with_video,
    create_user_message_with_video_url,
)
from vi.inference.predictors.hf import HFPredictor
from vi.inference.predictors.registry import PredictorRegistry
from vi.inference.types import SourceType
from vi.inference.utils.media_utils import extract_video_metadata
from vi.inference.utils.module_import import check_imports
from vi.inference.utils.path_utils import is_video_data_uri, is_video_file, is_video_url

check_imports(
    packages=["torch", "xgrammar"],
)

# Import after dependency check (intentional)
from vi.inference.config.qwenvl import QwenVLGenerationConfig  # noqa: E402


@PredictorRegistry.register(
    predictor_key="qwen25vl",
    loader_types=["Qwen25VLLoader"],
)
@PredictorRegistry.register(
    predictor_key="qwen3vl",
    loader_types=["Qwen3VLLoader"],
)
@PredictorRegistry.register(
    predictor_key="cosmosreason1",
    loader_types=["Qwen25VLLoader"],
)
@PredictorRegistry.register(
    predictor_key="cosmosreason2",
    loader_types=["Qwen3VLLoader"],
)
@PredictorRegistry.register(
    predictor_key="internvl35",
    loader_types=["Qwen25VLLoader"],
)
@PredictorRegistry.register(
    predictor_key="llavanext",
    loader_types=["Qwen25VLLoader"],
)
class QwenVLPredictor(HFPredictor[QwenVLGenerationConfig]):
    """Predictor for Qwen-VL compatible vision-language models.

    Handles preprocessing, inference, and output parsing for vision-language
    tasks. Works with Datature Vi fine-tuned models that include structured
    output generation via xgrammar.

    This predictor supports multiple vision-language model architectures that
    share the same preprocessing pipeline, making them fully compatible and
    interchangeable.

    Supported model architectures (all support both image and video inputs):
        - Qwen2.5-VL (Qwen2_5_VLForConditionalGeneration)
        - Qwen3-VL (Qwen3VLForConditionalGeneration)
        - InternVL 3.5 (InternVLForConditionalGeneration)
        - Cosmos Reason1 (Qwen2_5_VLForConditionalGeneration)
        - Cosmos Reason2 (Qwen3VLForConditionalGeneration)
        - LLaVA-NeXT (LlavaNextForConditionalGeneration)

    Supported task types:
        - Visual Question Answering (VQA): User prompt is required in the form of a question
        - Phrase Grounding: User prompt is optional (uses default prompt if not provided)
        - Freeform: Image captioning
        - Video Freeform: Video captioning (input modality automatically detected by file extension)

    Example:
        ```python
        from vi.inference.loaders import ViLoader
        from vi.inference.predictors import QwenVLPredictor

        # Load Qwen2.5-VL model and create predictor
        loader = ViLoader.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")
        predictor = QwenVLPredictor(loader)

        # Load Qwen3-VL model and create predictor
        loader = ViLoader.from_pretrained("Qwen/Qwen3-VL-7B-Instruct")
        predictor = QwenVLPredictor(loader)

        # Load InternVL 3.5 model and create predictor
        loader = ViLoader.from_pretrained("OpenGVLab/InternVL3.5-8B")
        predictor = QwenVLPredictor(loader)

        # Run inference
        result = predictor(
            source="image.jpg", user_prompt="What objects are in this image?"
        )
        print(result)
        ```

    Note:
        This predictor is designed for Datature Vi fine-tuned models with
        processor and xgrammar compiler. For pretrained models without these
        components, you'll need to implement a custom predictor.

        The generation config class used is QwenVLGenerationConfig, which serves
        as the base for all Qwen-VL versions (2.5, 3.0, etc.) due to shared
        tokenizer design and identical token IDs.

    """

    def get_generation_config_class(self) -> type[QwenVLGenerationConfig]:
        """Return the generation config class for Qwen-VL compatible models.

        Returns:
            QwenVLGenerationConfig class.

        """
        return QwenVLGenerationConfig

    def _validate_processor(self) -> None:
        """Validate that loader has a processor.

        Uses wrap_configuration_error for better error messages.

        Raises:
            AttributeError: If processor is missing or None.

        """
        if not hasattr(self._loader, "processor") or self._loader.processor is None:
            raise wrap_configuration_error(
                AttributeError(
                    "Predictor requires a loader with a processor. "
                    "The provided loader may not be properly configured."
                ),
                component="processor",
            )

    def _validate_compiler(self) -> None:
        """Validate that loader has a compiler.

        Uses wrap_configuration_error for better error messages.

        Raises:
            AttributeError: If compiler is missing or None.

        """
        if not hasattr(self._loader, "compiler") or self._loader.compiler is None:
            raise wrap_configuration_error(
                AttributeError(
                    "Predictor requires a loader with an xgrammar compiler. "
                    "This predictor is designed for Datature Vi fine-tuned models."
                ),
                component="compiler",
            )

    def _validate_metadata(self) -> None:
        """Validate that loader metadata has required fields.

        Uses wrap_configuration_error for better error messages.

        Raises:
            KeyError: If required metadata fields are missing.

        """
        if "task_type" not in self._loader.metadata:
            raise wrap_configuration_error(
                KeyError(
                    "Loader metadata missing 'task_type'. This predictor requires a "
                    "Datature Vi fine-tuned model with task type configuration."
                ),
                component="metadata",
            )

    def _prepare_inputs(
        self,
        source: str,
        source_type: SourceType,
        effective_prompt: str,
        system_prompt: str,
        fps: float,
    ) -> dict:
        """Prepare model inputs from image or video source and prompts.

        Uses the processor's apply_chat_template method for preprocessing.
        Automatically detects whether the source is an image or video based
        on file extension or source type.

        All models using this predictor (Qwen2.5-VL, Qwen3-VL, Cosmos Reason1,
        Cosmos Reason2, InternVL 3.5, LLaVA-NeXT) support both image and video inputs.

        For video inputs, extracts metadata (duration, fps, total frames) to enable
        proper fps-based frame sampling.

        Args:
            source: Image/video source - can be file path, URL, or data URI.
            source_type: Type of source - "path", "url", or "data_uri".
            effective_prompt: User prompt to use.
            system_prompt: System prompt (may include CoT suffix).
            fps: Frames per second for video sampling (only used for video inputs).

        Returns:
            Dictionary of model inputs ready for generation.

        """
        # Detect if source is a video based on source type
        if source_type == SourceType.PATH:
            video_input = is_video_file(source)
        elif source_type == SourceType.URL:
            video_input = is_video_url(source)
        else:  # SourceType.DATA_URI
            video_input = is_video_data_uri(source)

        # Extract video metadata before creating messages (needed for data URIs)
        video_metadata = None
        if video_input:
            # Extract metadata for file paths and data URIs
            # For URLs, the processor can extract metadata when it downloads
            if source_type in (SourceType.PATH, SourceType.DATA_URI):
                video_metadata = extract_video_metadata(source)

        # Create appropriate message type based on input type and source type
        if source_type in (SourceType.URL, SourceType.DATA_URI):
            # URLs and data URIs use the URL message functions
            user_message = (
                create_user_message_with_video_url(source, effective_prompt)
                if video_input
                else create_user_message_with_image_url(source, effective_prompt)
            )
        else:
            # File paths use the path message functions
            user_message = (
                create_user_message_with_video(source, effective_prompt)
                if video_input
                else create_user_message_with_image(source, effective_prompt)
            )

        messages = [
            create_system_message(system_prompt),
            user_message,
        ]

        # Prepare processor arguments
        processor_kwargs = {
            "tokenize": True,
            "add_generation_prompt": True,
            "return_dict": True,
            "return_tensors": "pt",
        }

        # For video inputs, pass fps and metadata (if extracted)
        if video_input:
            processor_kwargs["fps"] = fps
            if video_metadata is not None:
                processor_kwargs["video_metadata"] = video_metadata

        # Use processor.apply_chat_template with appropriate parameters
        inputs = self._loader.processor.apply_chat_template(
            messages, **processor_kwargs
        )
        inputs = inputs.to(self._loader.model.device)

        return inputs
