#!/usr/bin/env python

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

@File    :   vi_loader.py
@Author  :   Wei Loon Cheng
@Contact :   developers@datature.io
@License :   Apache License 2.0
@Desc    :   Datature Vi SDK auto loader module.
"""

import json
from pathlib import Path
from typing import Any

import msgspec
from transformers import AutoConfig
from vi.api.resources.flows.responses import FlowSpec
from vi.api.resources.models.results import ModelDownloadResult
from vi.inference.loaders.base_loader import BaseLoader
from vi.inference.loaders.registry import LoaderRegistry


class _ViLoaderFactory:
    """Factory for automatically loading models based on architecture detection.

    Detects model architecture from configuration files and automatically selects
    the appropriate loader. Supports HuggingFace models, local models, and
    Datature Vi fine-tuned models.

    !!! note "Recommended Usage"
        For most use cases, prefer using `ViModel` which handles client initialization
        and model downloading automatically:

        ```python
        from vi.inference import ViModel

        # Recommended: Use ViModel for Datature Vi models
        model = ViModel(
            secret_key="your-secret-key",
            organization_id="your-organization-id",
            run_id="your-run-id",
        )
        ```

    Detection Strategy:
        1. Check run.json for Datature Vi models (if present)
        2. Check config.json for model_type field
        3. Check config.json for architectures field
        4. Attempt pattern matching on model path/name
        5. Raise error if detection fails

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

        # Load Datature model from metadata
        model_meta = client.get_model(run_id="my-run")
        loader = ViLoader(model_meta=model_meta)

        # Load from local path (auto-detect)
        loader = ViLoader(pretrained_model_name_or_path="./my_model")

        # Load from HuggingFace (auto-detect)
        loader = ViLoader(pretrained_model_name_or_path="Qwen/Qwen2.5-VL-7B-Instruct")

        # Load with explicit loader selection
        loader = ViLoader(
            pretrained_model_name_or_path="Qwen/Qwen2.5-VL-7B-Instruct",
            loader_key="qwen25vl",
        )

        # Alternative explicit methods (still available)
        loader = ViLoader.from_pretrained("Qwen/Qwen2.5-VL-7B-Instruct")
        loader = ViLoader.from_model_meta(model_meta)
        ```

    Note:
        This factory supports both HuggingFace-compatible models and custom
        non-HuggingFace models. For non-HuggingFace models, implement a custom
        loader and register it with LoaderRegistry.

    """

    def __call__(
        self,
        model_meta: ModelDownloadResult | None = None,
        pretrained_model_name_or_path: str | None = None,
        loader_key: str | None = None,
        **kwargs: Any,
    ) -> BaseLoader:
        """Create a loader instance with automatic architecture detection.

        This is the primary way to create loaders. Supports both HuggingFace models
        and Datature Vi fine-tuned models with automatic architecture detection.

        Args:
            model_meta: Model metadata containing paths and configuration.
                Use this for Datature Vi fine-tuned models.
            pretrained_model_name_or_path: Model identifier or path. Can be:
                - HuggingFace model name (e.g., "Qwen/Qwen2.5-VL-7B-Instruct")
                - Local directory path containing model files
                - Any path with standard HuggingFace structure (config.json, etc.)
            loader_key: Optional explicit loader to use (e.g., "qwen25vl").
                If provided, skips automatic detection.
            **kwargs: Additional arguments passed to the loader:
                - attn_implementation: Attention implementation ("eager",
                  "flash_attention_2", "sdpa")
                - device_map: Device mapping ("auto", "cpu", "cuda", or dict)
                - low_cpu_mem_usage: Use low CPU memory mode (bool)
                - trust_remote_code: Trust remote code (bool)
                - max_threads: Max threads for operations (int)

        Returns:
            Loaded model instance ready for inference.

        Raises:
            ValueError: If neither model_meta nor pretrained_model_name_or_path is provided,
                or if model architecture cannot be detected or is not supported.
            FileNotFoundError: If local model path doesn't exist.
            ImportError: If required dependencies for the model are not installed.

        Example:
            ```python
            # Load from HuggingFace (auto-detect)
            loader = ViLoader(
                pretrained_model_name_or_path="Qwen/Qwen2.5-VL-7B-Instruct"
            )

            # Load from local path (auto-detect)
            loader = ViLoader(pretrained_model_name_or_path="./my_model")

            # Load with explicit loader selection
            loader = ViLoader(
                pretrained_model_name_or_path="Qwen/Qwen2.5-VL-7B-Instruct",
                loader_key="qwen25vl",
                trust_remote_code=True,
                device_map="auto",
            )

            # Load Datature model from metadata
            model_meta = client.get_model(run_id="my-run")
            loader = ViLoader(model_meta=model_meta)

            # Access loaded components
            print(f"Model: {loader.model}")
            print(f"Metadata: {loader.metadata}")
            ```

        Note:
            For HuggingFace models that aren't cached locally, the first call
            will download the model. Subsequent calls will use the cached version.

        """
        if model_meta and not pretrained_model_name_or_path:
            return self.from_model_meta(
                model_meta=model_meta,
                loader_key=loader_key,
                **kwargs,
            )

        if pretrained_model_name_or_path and not model_meta:
            return self.from_pretrained(
                pretrained_model_name_or_path=pretrained_model_name_or_path,
                loader_key=loader_key,
                **kwargs,
            )

        raise ValueError(
            "Either model_meta or pretrained_model_name_or_path must be provided"
        )

    def from_pretrained(
        self,
        pretrained_model_name_or_path: str,
        loader_key: str | None = None,
        **kwargs: Any,
    ) -> BaseLoader:
        """Load a model from pretrained path with automatic architecture detection.

        Args:
            pretrained_model_name_or_path: Model identifier or path. Can be:
                - HuggingFace model name (e.g., "Qwen/Qwen2.5-VL-7B-Instruct")
                - Local directory path containing model files
                - Any path with standard HuggingFace structure (config.json, etc.)
            loader_key: Optional explicit loader to use (e.g., "qwen25vl").
                If provided, skips automatic detection.
            **kwargs: Additional arguments passed to the loader:
                - attn_implementation: Attention implementation
                  ("eager", "flash_attention_2", "sdpa")
                - device_map: Device mapping ("auto", "cpu", "cuda", or dict)
                - low_cpu_mem_usage: Use low CPU memory mode (bool)
                - trust_remote_code: Trust remote code (bool)
                - max_threads: Max threads for operations (int)

        Returns:
            Loaded model instance ready for inference.

        Raises:
            ValueError: If model architecture cannot be detected or is not supported.
            FileNotFoundError: If local model path doesn't exist.
            ImportError: If required dependencies for the model are not installed.

        Example:
            ```python
            # Auto-detect from HuggingFace
            loader = ViLoader.from_pretrained(
                "Qwen/Qwen2.5-VL-7B-Instruct", trust_remote_code=True, device_map="auto"
            )

            # Auto-detect from local path
            loader = ViLoader.from_pretrained(
                "./models/my_finetuned_qwen", attn_implementation="flash_attention_2"
            )

            # Explicit loader selection
            loader = ViLoader.from_pretrained(
                "Qwen/Qwen2.5-VL-7B-Instruct",
                loader_key="qwen25vl",
                low_cpu_mem_usage=True,
            )

            # Access loaded components
            print(f"Model: {loader.model}")
            print(f"Metadata: {loader.metadata}")
            ```

        Note:
            For HuggingFace models that aren't cached locally, the first call
            will download the model. Subsequent calls will use the cached version.

        """
        model_meta = ModelDownloadResult(
            model_path=pretrained_model_name_or_path,
            adapter_path=None,
            run_config_path=None,
        )
        return self.from_model_meta(
            model_meta=model_meta, loader_key=loader_key, **kwargs
        )

    def from_model_meta(
        self,
        model_meta: ModelDownloadResult,
        loader_key: str | None = None,
        **kwargs: Any,
    ) -> BaseLoader:
        """Load a model from metadata with automatic architecture detection.

        Args:
            model_meta: Model metadata containing:
                - model_path: Path to model directory or HuggingFace identifier
                - adapter_path: Optional path to PEFT adapter
                - run_config_path: Optional path to Datature run.json
            loader_key: Optional explicit loader to use (e.g., "qwen25vl").
                If provided, skips automatic detection.
            **kwargs: Additional arguments passed to the loader.

        Returns:
            Loaded model instance ready for inference.

        Raises:
            ValueError: If model architecture cannot be detected or is not supported.
            FileNotFoundError: If local model files don't exist.
            ImportError: If required dependencies for the model are not installed.

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

            # Load Datature fine-tuned model
            model_meta = ModelDownloadResult(
                model_path="./models/run_123/model_full",
                adapter_path="./models/run_123/model_adapter",
                run_config_path="./models/run_123/model_full/run.json",
            )
            loader = ViLoader.from_model_meta(model_meta)

            # Load pretrained model
            model_meta = ModelDownloadResult(
                model_path="Qwen/Qwen2.5-VL-7B-Instruct",
                adapter_path=None,
                run_config_path=None,
            )
            loader = ViLoader.from_model_meta(model_meta, trust_remote_code=True)
            ```

        """
        # Auto-detect loader if not explicitly provided
        if loader_key is None:
            loader_key = self._detect_loader(model_meta)

        # Get loader class from registry
        loader_class = LoaderRegistry.get_loader(loader_key=loader_key)

        # Instantiate and return loader
        # Note: overwrite is used in the download phase, not passed to loader
        return loader_class(model_meta, **kwargs)

    def _detect_loader(self, model_meta: ModelDownloadResult) -> str:
        """Detect appropriate loader from model configuration.

        Detection order:
            1. Check for Datature run.json (for fine-tuned models)
            2. Download/read config.json and check model_type field
            3. Download/read config.json and check architectures field

        Args:
            model_meta: Model metadata to analyze.

        Returns:
            Loader key (e.g., "qwen25vl") for the detected model architecture.

        Raises:
            ValueError: If model architecture cannot be detected or config is missing.

        """
        # Strategy 1: Check for Datature run.json
        if model_meta.run_config_path and Path(model_meta.run_config_path).exists():
            try:
                loader_key = self._detect_from_datature_config(model_meta)
                if loader_key:
                    return loader_key
            except Exception:  # noqa: S110
                pass  # Fall through to config-based detection

        # Strategy 2: Load config.json (from HuggingFace or local)
        config_path = Path(model_meta.model_path) / "config.json"
        config = None

        # Try to download from HuggingFace if not cached locally
        if not config_path.exists():
            try:
                hf_config = AutoConfig.from_pretrained(
                    model_meta.model_path, trust_remote_code=False
                )
                config = hf_config.to_dict()
            except Exception as e:
                raise ValueError(
                    f"Could not find config.json at {config_path} and could not "
                    f"download config from HuggingFace for '{model_meta.model_path}'. "
                    f"Error: {e}"
                ) from e

        # Load from local file if not loaded from HuggingFace
        if config is None:
            try:
                with open(config_path, encoding="utf-8") as f:
                    config = json.loads(f.read())
            except (json.JSONDecodeError, OSError) as e:
                raise ValueError(
                    f"Failed to parse config.json at {config_path}: {e}"
                ) from e

        # Ensure loaders are discovered before checking registry
        LoaderRegistry._discover_architectures()

        # Strategy 2a: Check model_type field
        model_type = config.get("model_type")
        if model_type and LoaderRegistry.is_registered(model_type=model_type):
            loader_class = LoaderRegistry.get_loader(model_type=model_type)
            loader_key = LoaderRegistry.get_loader_key(loader_class)
            if loader_key:
                return loader_key

        # Strategy 2b: Check architectures field
        architectures = config.get("architectures", [])
        for arch in architectures:
            if LoaderRegistry.is_registered(architecture=arch):
                loader_class = LoaderRegistry.get_loader(architecture=arch)
                loader_key = LoaderRegistry.get_loader_key(loader_class)
                if loader_key:
                    return loader_key

        # If we get here, the model architecture is not supported
        raise ValueError(
            f"Model architecture not supported for '{model_meta.model_path}'. "
            f"Config has model_type='{model_type}', architectures={architectures}. "
            f"Supported model types: {LoaderRegistry.list_model_types()}. "
            f"Supported architectures: {LoaderRegistry.list_architectures()}."
        )

    def _detect_from_datature_config(
        self, model_meta: ModelDownloadResult
    ) -> str | None:
        """Detect loader from Datature run.json configuration.

        Args:
            model_meta: Model metadata with run_config_path.

        Returns:
            Loader key if detected, None otherwise.

        """
        if not model_meta.run_config_path:
            return None

        try:
            with open(model_meta.run_config_path, encoding="utf-8") as f:
                config = json.loads(f.read())

            flow_spec = msgspec.convert(config["spec"]["flow"], type=FlowSpec)

            model_block = next(
                (block for block in flow_spec.blocks if "model" in block.block),
                None,
            )
            if model_block:
                return model_block.settings["architecture"]["name"]

        except Exception:  # noqa: S110
            pass

        return None

    def list_available_loaders(self) -> list[str]:
        """List all available loader keys.

        Returns:
            List of registered loader keys.

        Example:
            ```python
            loaders = ViLoader.list_available_loaders()
            print(f"Available loaders: {loaders}")
            # Output: Available loaders: ['qwen25vl', ...]
            ```

        """
        return LoaderRegistry.list_loaders()

    def list_supported_model_types(self) -> list[str]:
        """List all supported model types from config.json.

        Returns:
            List of model_type values that can be automatically detected.

        Example:
            ```python
            types = ViLoader.list_supported_model_types()
            print(f"Supported model types: {types}")
            # Output: Supported model types: ['qwen2_5_vl', ...]
            ```

        """
        return LoaderRegistry.list_model_types()

    def list_supported_architectures(self) -> list[str]:
        """List all supported architecture names from config.json.

        Returns:
            List of architecture names that can be automatically detected.

        Example:
            ```python
            archs = ViLoader.list_supported_architectures()
            print(f"Supported architectures: {archs}")
            # Output: ['Qwen2_5_VLForConditionalGeneration', ...]
            ```

        """
        return LoaderRegistry.list_architectures()


ViLoader: _ViLoaderFactory = _ViLoaderFactory()
