from __future__ import annotations

import copy
import warnings
from collections.abc import Callable
from typing import TYPE_CHECKING, Optional, Union

if TYPE_CHECKING:
    from typing_extensions import TypeAlias

from urllib.parse import urlparse

from ramalama.arg_types import StoreArgType
from ramalama.chat_providers.api_providers import get_chat_provider
from ramalama.common import rm_until_substring
from ramalama.config import ActiveConfig
from ramalama.path_utils import file_uri_to_path
from ramalama.transports.api import APITransport
from ramalama.transports.base import MODEL_TYPES, Transport
from ramalama.transports.huggingface import Huggingface
from ramalama.transports.modelscope import ModelScope
from ramalama.transports.oci.oci import OCI
from ramalama.transports.ollama import Ollama
from ramalama.transports.rlcr import RamalamaContainerRegistry
from ramalama.transports.url import URL

CLASS_MODEL_TYPES: TypeAlias = Union[Huggingface, Ollama, OCI, URL, ModelScope, RamalamaContainerRegistry, APITransport]

_ollama_default_warned = False


class TransportFactory:
    def __init__(
        self,
        model: str,
        args: StoreArgType,
        transport: Optional[str] = None,
        ignore_stderr: bool = False,
    ):

        self.model = model
        self.store_path = args.store
        self._transport_is_implicit = transport is None and not model.startswith("ollama://")
        self.transport = transport if transport is not None else "ollama"
        self.engine = args.engine
        self.ignore_stderr = ignore_stderr
        self.container = args.container

        model_cls, _create = self.detect_model_model_type()

        self.model_cls = model_cls
        self._create = _create

        self.pruned_model = self.prune_model_input()
        self.draft_model: Optional[Transport] = None

        model_draft = getattr(args, "model_draft", None)
        if model_draft:
            dm_args = copy.deepcopy(args)
            dm_args.model_draft = None  # type: ignore
            draft_model = TransportFactory(model_draft, dm_args, ignore_stderr=True).create()
            if not isinstance(draft_model, Transport):
                raise ValueError("Draft models must be local transports; hosted API transports are not supported.")
            self.draft_model = draft_model

    def detect_model_model_type(self) -> tuple[type[CLASS_MODEL_TYPES], Callable[[], CLASS_MODEL_TYPES]]:
        if self.model.startswith(("huggingface://", "hf://", "hf.co/")):
            return Huggingface, self.create_huggingface
        if self.model.startswith(("modelscope://", "ms://")):
            return ModelScope, self.create_modelscope
        if self.model.startswith("ollama://"):
            self._warn_ollama_deprecated()
            return Ollama, self.create_ollama
        if self.model.startswith(("oci://", "docker://")):
            return OCI, self.create_oci
        if self.model.startswith("rlcr://"):
            return RamalamaContainerRegistry, self.create_rlcr
        if self.model.startswith(("http://", "https://", "file:")):
            return URL, self.create_url
        if self.model.startswith(("openai://")):
            return APITransport, self.create_api_transport
        if self._transport_is_implicit and "/" in self.model:
            return Huggingface, self.create_huggingface
        if self.transport == "huggingface":
            return Huggingface, self.create_huggingface
        if self.transport == "modelscope":
            return ModelScope, self.create_modelscope
        if self.transport == "ollama":
            self._warn_ollama_deprecated()
            return Ollama, self.create_ollama
        if self.transport == "rlcr":
            return RamalamaContainerRegistry, self.create_rlcr
        if self.transport == "oci":
            return OCI, self.create_oci
        raise KeyError(f'transport "{self.transport}" not supported. Must be oci, huggingface, modelscope, or ollama.')

    def _warn_ollama_deprecated(self) -> None:
        global _ollama_default_warned
        if _ollama_default_warned:
            return
        _ollama_default_warned = True
        msg = ""
        if self._transport_is_implicit:
            msg = f"Defaulting to Ollama transport for '{self.model}'.\n"
        warnings.warn(
            f"{msg}Note: Ollama models are no longer compatible with llama.cpp. "
            "Support for Ollama models will be removed in a future release. "
            "Use the Hugging Face <org>/<model> shorthand format, "
            "or specify a transport explicitly, e.g. huggingface:// or oci://.",
            FutureWarning,
            stacklevel=3,
        )

    def prune_model_input(self) -> str:

        if self.model_cls == URL and urlparse(self.model).scheme == "file":
            return file_uri_to_path(self.model)

        # remove protocol from model input
        pruned_model_input = rm_until_substring(self.model, "://")

        if self.model_cls == Huggingface:
            pruned_model_input = rm_until_substring(pruned_model_input, "hf.co/")
        elif self.model_cls == ModelScope:
            pruned_model_input = rm_until_substring(pruned_model_input, "modelscope.cn/")
        elif self.model_cls == Ollama:
            pruned_model_input = rm_until_substring(pruned_model_input, "library/")

        return pruned_model_input

    def validate_oci_model_input(self):
        if self.model.startswith("oci://") or self.model.startswith("docker://") or self.model.startswith("rlcr://"):
            return

        for t in MODEL_TYPES:
            if self.model.startswith(t + "://"):
                raise ValueError(f"{self.model} invalid: Only OCI Model types supported")

    def create(self) -> CLASS_MODEL_TYPES:
        return self._create()

    def create_huggingface(self) -> Huggingface:
        model = Huggingface(self.pruned_model, self.store_path)
        model.draft_model = self.draft_model
        return model

    def create_modelscope(self) -> ModelScope:
        model = ModelScope(self.pruned_model, self.store_path)
        model.draft_model = self.draft_model
        return model

    def create_ollama(self) -> Ollama:
        model = Ollama(self.pruned_model, self.store_path)
        model.draft_model = self.draft_model
        return model

    def create_rlcr(self) -> RamalamaContainerRegistry:
        if not self.container:
            raise ValueError("OCI containers cannot be used with the --nocontainer option.")

        if self.engine is None:
            raise ValueError("Constructing an OCI model factory requires an engine value")

        self.validate_oci_model_input()
        model = RamalamaContainerRegistry(
            model=self.pruned_model,
            model_store_path=self.store_path,
            conman=self.engine,
            ignore_stderr=self.ignore_stderr,
        )
        model.draft_model = self.draft_model
        return model

    def create_oci(self) -> OCI:
        if not self.container:
            raise ValueError("OCI containers cannot be used with the --nocontainer option.")

        if self.engine is None:
            raise ValueError("Constructing an OCI model factory requires an engine value")

        self.validate_oci_model_input()

        model = OCI(self.pruned_model, self.store_path, self.engine, self.ignore_stderr)
        model.draft_model = self.draft_model
        return model

    def create_url(self) -> URL:
        model = URL(self.pruned_model, self.store_path, urlparse(self.model).scheme)
        model.draft_model = self.draft_model
        return model

    def create_api_transport(self) -> APITransport:
        scheme = self.model.split("://", 1)[0]
        return APITransport(self.pruned_model, provider=get_chat_provider(scheme))


def New(name, args, transport: Optional[str] = None) -> CLASS_MODEL_TYPES:
    if transport is None:
        config = ActiveConfig()
        transport = config.transport if config.is_set("transport") else None
    return TransportFactory(name, args, transport=transport).create()
