#!/usr/bin/env python

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

@File    :   qwen3vl.py
@Author  :   Wei Loon Cheng
@Contact :   developers@datature.io
@License :   Apache License 2.0
@Desc    :   Datature Vi SDK Qwen3-VL loader module.
"""

from vi.inference.loaders.hf import HFLoader
from vi.inference.loaders.registry import LoaderRegistry
from vi.inference.utils.module_import import check_imports

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

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


@LoaderRegistry.register(
    loader_key="qwen3vl",
    model_types=["qwen3_vl"],
    architectures=["Qwen3VLForConditionalGeneration"],
)
@LoaderRegistry.register(
    loader_key="cosmosreason2",
    model_types=["qwen3_vl"],
    architectures=["Qwen3VLForConditionalGeneration"],
)
class Qwen3VLLoader(HFLoader):
    """Loader for Qwen3-VL vision-language models.

    Handles loading of Qwen3-VL based vision-language models from HuggingFace Hub,
    local paths, or Datature Vi fine-tuned models. Supports optional PEFT adapters,
    quantization, and structured output generation via xgrammar.

    Qwen3-VL is the next generation of Qwen vision-language models with improved
    reasoning capabilities and multimodal understanding.

    Supported model architectures:
        - Qwen3-VL (Qwen3VLForConditionalGeneration)
        - NVIDIA Cosmos Reason2 (Qwen3VLForConditionalGeneration)

    Supported configurations:
        - Pretrained models from HuggingFace:
            - "Qwen/Qwen3-VL-*" models
            - "nvidia/Cosmos-Reason2-2B"
        - Fine-tuned models from Datature Vi platform
        - Local model directories with standard HuggingFace structure
        - 4-bit and 8-bit quantization for memory efficiency
        - PEFT adapters (LoRA, QLoRA)
        - Structured output generation (requires xgrammar)

    Example:
        ```python
        from vi.inference.loaders import Qwen3VLLoader
        from vi.api.resources.models.results import ModelDownloadResult

        # Load NVIDIA Cosmos Reason2-2B from HuggingFace
        model_meta = ModelDownloadResult(
            model_path="nvidia/Cosmos-Reason2-2B",
            adapter_path=None,
            run_config_path=None,
        )
        loader = Qwen3VLLoader(
            model_meta=model_meta,
            trust_remote_code=True,
        )

        # Or use ViModel for automatic detection
        from vi.inference import ViModel

        model = ViModel(
            pretrained_model_name_or_path="nvidia/Cosmos-Reason2-2B",
            loader_key="qwen3vl",  # or "cosmosreason2"
            trust_remote_code=True,
        )

        # Access components
        print(f"Model: {loader.model}")
        print(f"Processor: {loader.processor}")
        print(f"Compiler: {loader.compiler}")
        ```

    Note:
        Qwen3-VL models use the same processor interface as Qwen2.5-VL,
        making migration easy.

    """

    _generation_config_class = QwenVLGenerationConfig
    _model_class = AutoModelForImageTextToText
    _processor_class = AutoProcessor
