import entrypoints
import warnings
from abc import ABCMeta, abstractmethod

from segmind.exceptions import MlflowException
from segmind.protos.errorcodes_pb2 import INVALID_PARAMETER_VALUE
from segmind.utils.uri import get_uri_scheme


class UnsupportedModelRegistryStoreURIException(MlflowException):
    """Exception thrown when building a model registry store with an
    unsupported URI."""

    def __init__(self, unsupported_uri, supported_uri_schemes):
        message = f"Unsupported URI '{unsupported_uri}' for model registry \
        store. Supported schemes are: {supported_uri_schemes}"
        super(UnsupportedModelRegistryStoreURIException, self).__init__(
            message, error_code=INVALID_PARAMETER_VALUE)
        self.supported_uri_schemes = supported_uri_schemes


class StoreRegistry:
    """Abstract class defining a scheme-based registry for store
    implementations.

    This class allows the registration of a function or class to provide an
    implementation for a given scheme of `store_uri` through the `register`
    methods. Implementations declared though the entrypoints can be
    automatically registered through the `register_entrypoints` method.

    When instantiating a store through the `get_store` method, the scheme of
    the store URI provided (or inferred from environment) will be used to
    select which implementation to instantiate, which will be called with same
    arguments passed to the `get_store` method.
    """

    __metaclass__ = ABCMeta

    @abstractmethod
    def __init__(self, group_name):
        self._registry = {}
        self.group_name = group_name

    def register(self, scheme, store_builder):
        self._registry[scheme] = store_builder

    def register_entrypoints(self):
        """Register tracking stores provided by other packages."""
        for entrypoint in entrypoints.get_group_all(self.group_name):
            try:
                self.register(entrypoint.name, entrypoint.load())
            except (AttributeError, ImportError) as exc:
                warnings.warn(
                    f'Failure attempting to register store for scheme \
                    "{entrypoint.name}": {str(exc)}',
                    stacklevel=2)

    def get_store_builder(self, store_uri):
        """Get a store from the registry based on the scheme of store_uri.

        Args:
        store_uri: The store URI. If None, it will be inferred from the
        environment. This URI is used to select which tracking store
        implementation to instantiate and is passed to the constructor of the
        implementation.
        :return: A function that returns an instance of
                 ``segmind_track.store.{tracking|model_registry}.AbstractStore
                 `` that fulfills the store URI requirements.
        """
        scheme = store_uri if store_uri == 'databricks' else get_uri_scheme(
            store_uri)

        try:
            store_builder = self._registry[scheme]
        except KeyError:
            raise UnsupportedModelRegistryStoreURIException(
                unsupported_uri=store_uri,
                supported_uri_schemes=list(self._registry.keys()))
        return store_builder
