from __future__ import annotations
from typing import Dict, Text, Any, Optional

from bf_nlu.engine.graph import GraphComponent, ExecutionContext
from bf_nlu.engine.storage.resource import Resource
from bf_nlu.engine.storage.storage import ModelStorage
from bf_nlu.shared.core.domain import Domain
from bf_nlu.shared.exceptions import InvalidConfigException
from bf_nlu.shared.importers.importer import TrainingDataImporter


class DomainProvider(GraphComponent):
    """Provides domain during training and inference time."""

    def __init__(
        self,
        model_storage: ModelStorage,
        resource: Resource,
        domain: Optional[Domain] = None,
    ) -> None:
        """Creates domain provider."""
        self._model_storage = model_storage
        self._resource = resource
        self._domain = domain

    @classmethod
    def create(
        cls,
        config: Dict[Text, Any],
        model_storage: ModelStorage,
        resource: Resource,
        execution_context: ExecutionContext,
    ) -> DomainProvider:
        """Creates component (see parent class for full docstring)."""
        return cls(model_storage, resource)

    @classmethod
    def load(
        cls,
        config: Dict[Text, Any],
        model_storage: ModelStorage,
        resource: Resource,
        execution_context: ExecutionContext,
        **kwargs: Any,
    ) -> DomainProvider:
        """Creates provider using a persisted version of itself."""
        with model_storage.read_from(resource) as resource_directory:
            domain = Domain.from_path(resource_directory)
        return cls(model_storage, resource, domain)

    def _persist(self, domain: Domain) -> None:
        """Persists domain to model storage."""
        with self._model_storage.write_to(self._resource) as resource_directory:
            domain.persist(resource_directory / "domain.yml")

    def provide_train(self, importer: TrainingDataImporter) -> Domain:
        """Provides domain from training data during training."""
        domain = importer.get_domain()
        self._persist(domain)
        return domain

    def provide_inference(self) -> Domain:
        """Provides the domain during inference."""
        if self._domain is None:
            # This can't really happen but if it happens then we fail early
            raise InvalidConfigException(
                "No domain was found. This is required for "
                "making model predictions. Please make sure to "
                "provide a valid domain during training."
            )
        return self._domain
