#!/usr/bin/env python

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

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

from vi.inference.messages import (
    create_system_message,
    create_user_message_with_image,
    create_user_message_with_image_url,
)
from vi.inference.predictors.hf import HFPredictor
from vi.inference.predictors.registry import PredictorRegistry
from vi.inference.task_types.consts import VIDEO_SUPPORTED_MODELS_STR
from vi.inference.types import SourceType
from vi.inference.utils.module_import import check_imports
from vi.inference.utils.path_utils import is_video_file

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

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


@PredictorRegistry.register(
    predictor_key="nvila",
    loader_types=["NVILALoader"],
)
class NVILAPredictor(HFPredictor[NVILAGenerationConfig]):
    """NVILA predictor for vision-language tasks.

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

    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)

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

        # Load model and create predictor
        loader = ViLoader.from_pretrained("Efficient-Large-Model/NVILA-8B")
        predictor = NVILAPredictor(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.

    """

    def get_generation_config_class(self) -> type[NVILAGenerationConfig]:
        """Return the generation config class for NVILA models.

        Returns:
            NVILAGenerationConfig class.

        """
        return NVILAGenerationConfig

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

        Uses the processor's apply_chat_template method for preprocessing.

        Args:
            source: Image 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 (not used for NVILA, image-only model).

        Returns:
            Dictionary of model inputs ready for generation.

        Raises:
            ValueError: If video input is provided (not supported by this model).

        """
        # Check if input is a video file (only for path type)
        if source_type == SourceType.PATH and is_video_file(source):
            raise ValueError(
                f"Video input is not supported by NVILA models. "
                f"This model only supports image inputs. "
                f"For video inference, please use one of these models: {VIDEO_SUPPORTED_MODELS_STR}."
            )

        # Create appropriate message type based on source type
        if source_type in (SourceType.URL, SourceType.DATA_URI):
            user_message = create_user_message_with_image_url(source, effective_prompt)
        else:
            user_message = create_user_message_with_image(source, effective_prompt)

        messages = [
            create_system_message(system_prompt),
            user_message,
        ]

        # Use processor.apply_chat_template with tokenize=True
        inputs = self._loader.processor.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            return_dict=True,
            return_tensors="pt",
        )
        inputs = inputs.to(self._loader.model.device)

        return inputs
