from __future__ import annotations

from collections import defaultdict
from typing import Any, Dict, Iterable, List, Set, Text, Type, cast

import rasa.shared.utils.io
from rasa.core.constants import POLICY_PRIORITY
from rasa.core.featurizers.precomputation import CoreFeaturizationInputConverter
from rasa.core.policies.flow_policy import FlowPolicy
from rasa.core.policies.memoization import AugmentedMemoizationPolicy, MemoizationPolicy
from rasa.core.policies.policy import Policy, SupportedData
from rasa.core.policies.rule_policy import RulePolicy
from rasa.engine.graph import ExecutionContext, GraphComponent, GraphSchema, SchemaNode
from rasa.engine.storage.resource import Resource
from rasa.engine.storage.storage import ModelStorage
from rasa.nlu.extractors.entity_synonyms import EntitySynonymMapper
from rasa.nlu.extractors.regex_entity_extractor import RegexEntityExtractor
from rasa.nlu.featurizers.featurizer import Featurizer
from rasa.nlu.featurizers.sparse_featurizer.regex_featurizer import RegexFeaturizer
from rasa.nlu.tokenizers.tokenizer import Tokenizer
from rasa.shared.constants import (
    DEFAULT_CONFIG_PATH,
    DOCS_URL_COMPONENTS,
    DOCS_URL_DEFAULT_ACTIONS,
    DOCS_URL_NLU_BASED_POLICIES,
    DOCS_URL_RULES,
)
from rasa.shared.core.constants import (
    ACTION_BACK_NAME,
    ACTION_RESTART_NAME,
    USER_INTENT_BACK,
    USER_INTENT_RESTART,
)
from rasa.shared.core.domain import Domain, InvalidDomain
from rasa.shared.core.training_data.structures import RuleStep, StoryGraph
from rasa.shared.exceptions import InvalidConfigException
from rasa.shared.importers.importer import TrainingDataImporter
from rasa.shared.nlu.training_data.training_data import TrainingData
from rasa.shared.utils.common import conditional_import

# Conditional imports for TensorFlow-dependent components
TEDPolicy, TED_POLICY_AVAILABLE = conditional_import(
    "rasa.core.policies.ted_policy", "TEDPolicy", check_installation_setup=True
)
UnexpecTEDIntentPolicy, UNEXPECTED_INTENT_POLICY_AVAILABLE = conditional_import(
    "rasa.core.policies.unexpected_intent_policy",
    "UnexpecTEDIntentPolicy",
    check_installation_setup=True,
)
DIETClassifier, DIET_CLASSIFIER_AVAILABLE = conditional_import(
    "rasa.nlu.classifiers.diet_classifier",
    "DIETClassifier",
    check_installation_setup=True,
)
ResponseSelector, RESPONSE_SELECTOR_AVAILABLE = conditional_import(
    "rasa.nlu.selectors.response_selector",
    "ResponseSelector",
    check_installation_setup=True,
)

# Conditional imports for nlu components requiring other dependencies than tensorflow
MitieEntityExtractor, MITIE_ENTITY_EXTRACTOR_AVAILABLE = conditional_import(
    "rasa.nlu.extractors.mitie_entity_extractor", "MitieEntityExtractor"
)
CRFEntityExtractor, CRF_ENTITY_EXTRACTOR_AVAILABLE = conditional_import(
    "rasa.nlu.extractors.crf_entity_extractor", "CRFEntityExtractor"
)
CRFEntityExtractorOptions, _ = conditional_import(
    "rasa.nlu.extractors.crf_entity_extractor", "CRFEntityExtractorOptions"
)

# TODO: Can we replace this with the registered types from the registry?
TRAINABLE_EXTRACTORS = []
if CRF_ENTITY_EXTRACTOR_AVAILABLE:
    TRAINABLE_EXTRACTORS.append(CRFEntityExtractor)
if MITIE_ENTITY_EXTRACTOR_AVAILABLE:
    TRAINABLE_EXTRACTORS.append(MitieEntityExtractor)
if DIET_CLASSIFIER_AVAILABLE:
    TRAINABLE_EXTRACTORS.append(DIETClassifier)
# TODO: replace these once the Recipe is merged (used in tests)
POLICY_CLASSSES = {MemoizationPolicy, RulePolicy}
if TED_POLICY_AVAILABLE:
    POLICY_CLASSSES.add(TEDPolicy)


def _types_to_str(types: Iterable[Type]) -> Text:
    """Returns a text containing the names of all given types.

    Args:
        types: some types
    Returns:
        text containing all type names
    """
    return ", ".join([type.__name__ for type in types])


class DefaultV1RecipeValidator(GraphComponent):
    """Validates a "DefaultV1" configuration against the training data and domain."""

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

    def __init__(self, graph_schema: GraphSchema) -> None:
        """Instantiates a new `ConfigValidator`.

        Args:
           graph_schema: a graph schema
        """
        self._graph_schema = graph_schema
        self._component_types = set(node.uses for node in graph_schema.nodes.values())
        self._policy_schema_nodes: List[SchemaNode] = [
            node
            for node in self._graph_schema.nodes.values()
            if issubclass(node.uses, Policy)
        ]

    def validate(self, importer: TrainingDataImporter) -> TrainingDataImporter:
        """Validates the current graph schema against the training data and domain.

        Args:
            importer: the training data importer which can also load the domain
        Raises:
            `InvalidConfigException` or `InvalidDomain` in case there is some mismatch
        """
        nlu_data = importer.get_nlu_data()
        self._validate_nlu(nlu_data)

        story_graph = importer.get_stories()
        domain = importer.get_domain()
        self._validate_core(story_graph, domain)
        return importer

    def _validate_nlu(self, training_data: TrainingData) -> None:
        """Validates whether the configuration matches the training data.

        Args:
           training_data: The training data for the NLU components.
        """
        training_data.validate()

        self._raise_if_more_than_one_tokenizer()
        self._raise_if_featurizers_are_not_compatible()
        self._warn_of_competing_extractors()
        self._warn_of_competition_with_regex_extractor(training_data=training_data)
        self._warn_if_some_training_data_is_unused(training_data=training_data)

    def _warn_if_some_training_data_is_unused(
        self, training_data: TrainingData
    ) -> None:
        """Validates that all training data will be consumed by some component.

        For example, if you specify response examples in your training data, but there
        is no `ResponseSelector` component in your configuration, then this method
        issues a warning.

        Args:
            training_data: The training data for the NLU components.
        """
        if (
            training_data.response_examples
            and ResponseSelector not in self._component_types
        ):
            rasa.shared.utils.io.raise_warning(
                f"You have defined training data with examples for training a response "
                f"selector, but your NLU configuration does not include a response "
                f"selector component. "
                f"To train a model on your response selector data, add a "
                f"'{ResponseSelector.__name__}' to your configuration.",
                docs=DOCS_URL_COMPONENTS,
            )

        if training_data.entity_examples and self._component_types.isdisjoint(
            TRAINABLE_EXTRACTORS
        ):
            rasa.shared.utils.io.raise_warning(
                f"You have defined training data consisting of entity examples, but "
                f"your NLU configuration does not include an entity extractor "
                f"trained on your training data. "
                f"To extract non-pretrained entities, add one of "
                f"{_types_to_str(TRAINABLE_EXTRACTORS)} to your configuration.",
                docs=DOCS_URL_COMPONENTS,
            )

        available_extractors = {DIETClassifier}
        if CRF_ENTITY_EXTRACTOR_AVAILABLE:
            available_extractors.add(CRFEntityExtractor)

        if training_data.entity_examples and self._component_types.isdisjoint(
            available_extractors
        ):
            if training_data.entity_roles_groups_used():
                crf_name = (
                    CRFEntityExtractor.__name__
                    if CRFEntityExtractor
                    else "CRFEntityExtractor"
                )
                rasa.shared.utils.io.raise_warning(
                    f"You have defined training data with entities that "
                    f"have roles/groups, but your NLU configuration does not "
                    f"include a '{DIETClassifier.__name__}' "
                    f"or a '{crf_name}'. "
                    f"To train entities that have roles/groups, "
                    f"add either '{DIETClassifier.__name__}' "
                    f"or '{crf_name}' to your "
                    f"configuration.",
                    docs=DOCS_URL_COMPONENTS,
                )

        if training_data.regex_features and self._component_types.isdisjoint(
            [RegexFeaturizer, RegexEntityExtractor]
        ):
            rasa.shared.utils.io.raise_warning(
                f"You have defined training data with regexes, but "
                f"your NLU configuration does not include a 'RegexFeaturizer' "
                f" or a "
                f"'RegexEntityExtractor'. To use regexes, include either a "
                f"'{RegexFeaturizer.__name__}' or a "
                f"'{RegexEntityExtractor.__name__}' "
                f"in your configuration.",
                docs=DOCS_URL_COMPONENTS,
            )

        if training_data.lookup_tables and self._component_types.isdisjoint(
            [RegexFeaturizer, RegexEntityExtractor]
        ):
            rasa.shared.utils.io.raise_warning(
                f"You have defined training data consisting of lookup tables, but "
                f"your NLU configuration does not include a featurizer "
                f"or an entity extractor using the lookup table."
                f"To use the lookup tables, include either a "
                f"'{RegexFeaturizer.__name__}' "
                f"or a '{RegexEntityExtractor.__name__}' "
                f"in your configuration.",
                docs=DOCS_URL_COMPONENTS,
            )

        if training_data.lookup_tables:
            lookup_extractors = [DIETClassifier]
            if CRF_ENTITY_EXTRACTOR_AVAILABLE:
                lookup_extractors.append(CRFEntityExtractor)

            if self._component_types.isdisjoint(lookup_extractors):
                crf_name = (
                    CRFEntityExtractor.__name__
                    if CRFEntityExtractor
                    else "CRFEntityExtractor"
                )
                rasa.shared.utils.io.raise_warning(
                    f"You have defined training data consisting of lookup tables, but "
                    f"your NLU configuration does not include any components "
                    f"that uses the features created from the lookup table. "
                    f"To make use of the features that are created with the "
                    f"help of the lookup tables, "
                    f"add a '{DIETClassifier.__name__}' or a "
                    f"'{crf_name}' "
                    f"with the 'pattern' feature "
                    f"to your configuration.",
                    docs=DOCS_URL_COMPONENTS,
                )

            elif (
                CRF_ENTITY_EXTRACTOR_AVAILABLE
                and CRFEntityExtractor in self._component_types
            ):
                crf_schema_nodes = [
                    schema_node
                    for schema_node in self._graph_schema.nodes.values()
                    if schema_node.uses == CRFEntityExtractor
                ]
                has_pattern_feature = (
                    any(
                        CRFEntityExtractorOptions.PATTERN in feature_list
                        for crf in crf_schema_nodes
                        for feature_list in crf.config.get("features", [])
                    )
                    if CRFEntityExtractorOptions
                    else False
                )

                if not has_pattern_feature:
                    crf_name = (
                        CRFEntityExtractor.__name__
                        if CRFEntityExtractor
                        else "CRFEntityExtractor"
                    )
                    rasa.shared.utils.io.raise_warning(
                        f"You have defined training data consisting of "
                        f"lookup tables, but your NLU configuration's "
                        f"'{crf_name}' "
                        f"does not include the "
                        f"'pattern' feature. To featurize lookup tables, "
                        f"add the 'pattern' feature to the "
                        f"'{crf_name}' "
                        "in your configuration.",
                        docs=DOCS_URL_COMPONENTS,
                    )

        if (
            training_data.entity_synonyms
            and EntitySynonymMapper not in self._component_types
        ):
            rasa.shared.utils.io.raise_warning(
                f"You have defined synonyms in your training data, but "
                f"your NLU configuration does not include an "
                f"'{EntitySynonymMapper.__name__}'. "
                f"To map synonyms, add an "
                f"'{EntitySynonymMapper.__name__}' to your "
                f"configuration.",
                docs=DOCS_URL_COMPONENTS,
            )

    def _raise_if_more_than_one_tokenizer(self) -> None:
        """Validates that only one tokenizer is present in the configuration.

        Note that the existence of a tokenizer and its position in the graph schema
        will be validated via the validation of required components during
        schema validation.

        Raises:
            `InvalidConfigException` in case there is more than one tokenizer
        """
        types_of_tokenizer_schema_nodes = [
            schema_node.uses
            for schema_node in self._graph_schema.nodes.values()
            if issubclass(schema_node.uses, Tokenizer) and schema_node.fn != "train"
        ]

        is_end_to_end = any(
            issubclass(schema_node.uses, CoreFeaturizationInputConverter)
            for schema_node in self._graph_schema.nodes.values()
        )

        allowed_number_of_tokenizers = 2 if is_end_to_end else 1
        if len(types_of_tokenizer_schema_nodes) > allowed_number_of_tokenizers:
            raise InvalidConfigException(
                f"The configuration configuration contains more than one tokenizer, "
                f"which is not possible at this time. You can only use one tokenizer. "
                f"The configuration contains the following tokenizers: "
                f"{_types_to_str(types_of_tokenizer_schema_nodes)}. "
            )

    def _warn_of_competing_extractors(self) -> None:
        """Warns the user when using competing extractors.

        Competing extractors are e.g. `CRFEntityExtractor` and `DIETClassifier`.
        Both of these look for the same entities based on the same training data
        leading to ambiguity in the results.
        """
        extractors_in_configuration: Set[Type[GraphComponent]] = (
            self._component_types.intersection(TRAINABLE_EXTRACTORS)
        )
        if len(extractors_in_configuration) > 1:
            rasa.shared.utils.io.raise_warning(
                f"You have defined multiple entity extractors that do the same job "
                f"in your configuration: "
                f"{_types_to_str(extractors_in_configuration)}. "
                f"This can lead to the same entity getting "
                f"extracted multiple times. Please read the documentation section "
                f"on entity extractors to make sure you understand the implications.",
                docs=f"{DOCS_URL_COMPONENTS}#entity-extractors",
            )

    def _warn_of_competition_with_regex_extractor(
        self, training_data: TrainingData
    ) -> None:
        """Warns when regex entity extractor is competing with a general one.

        This might be the case when the following conditions are all met:
        * You are using a general entity extractor and the `RegexEntityExtractor`
        * AND you have regex patterns for entity type A
        * AND you have annotated text examples for entity type A

        Args:
            training_data: The training data for the NLU components.
        """
        present_general_extractors = self._component_types.intersection(
            TRAINABLE_EXTRACTORS
        )
        has_general_extractors = len(present_general_extractors) > 0
        has_regex_extractor = RegexEntityExtractor in self._component_types

        regex_entity_types = {rf["name"] for rf in training_data.regex_features}
        overlap_between_types = training_data.entities.intersection(regex_entity_types)
        has_overlap = len(overlap_between_types) > 0

        if has_general_extractors and has_regex_extractor and has_overlap:
            rasa.shared.utils.io.raise_warning(
                f"You have an overlap between the "
                f"'{RegexEntityExtractor.__name__}' and the "
                f"statistical entity extractors "
                f"{_types_to_str(present_general_extractors)} "
                f"in your configuration. Specifically both types of extractors will "
                f"attempt to extract entities of the types "
                f"{', '.join(overlap_between_types)}. "
                f"This can lead to multiple "
                f"extraction of entities. Please read "
                f"'{RegexEntityExtractor.__name__}''s "
                f"documentation section to make sure you understand the "
                f"implications.",
                docs=f"{DOCS_URL_COMPONENTS}#regexentityextractor",
            )

    def _raise_if_featurizers_are_not_compatible(self) -> None:
        """Raises or warns if there are problems regarding the featurizers.

        Raises:
            `InvalidConfigException` in case the featurizers are not compatible
        """
        featurizers: List[SchemaNode] = [
            node
            for node_name, node in self._graph_schema.nodes.items()
            if issubclass(node.uses, Featurizer)
            # Featurizers are split in `train` and `process_training_data` -
            # we only need to look at the nodes which _add_ features.
            and node.fn == "process_training_data"
            # Tokenizers are re-used in the Core part of the graph when using End-to-End
            and not node_name.startswith("e2e")
        ]

        Featurizer.raise_if_featurizer_configs_are_not_compatible(
            [schema_node.config for schema_node in featurizers]
        )

    def _validate_core(self, story_graph: StoryGraph, domain: Domain) -> None:
        """Validates whether the configuration matches the training data.

        Args:
           story_graph: a story graph (core training data)
           domain: the domain
        """
        if not self._policy_schema_nodes and story_graph.story_steps:
            rasa.shared.utils.io.raise_warning(
                "Found data for training policies but no policy was configured.",
                docs=DOCS_URL_NLU_BASED_POLICIES,
            )
        if not self._policy_schema_nodes:
            return

        # self._raise_if_dm1_and_calm_policies_used_at_the_same_time()
        self._warn_if_no_policy_handles_default_intents()
        self._raise_if_domain_contains_form_names_but_no_rule_policy_given(domain)
        self._raise_if_a_rule_policy_is_incompatible_with_domain(domain)
        self._validate_policy_priorities()
        self._warn_if_rule_based_data_is_unused_or_missing(story_graph=story_graph)

    def _warn_if_no_policy_handles_default_intents(self) -> None:
        """Warns if there is no policy handling the default intents."""
        policies_handling_default_indents = (RulePolicy, FlowPolicy)

        contains_policy_handling_default_intents = any(
            issubclass(node.uses, policies_handling_default_indents)
            for node in self._policy_schema_nodes
        )
        if not contains_policy_handling_default_intents:
            rasa.shared.utils.io.raise_warning(
                f"'{RulePolicy.__name__}' is not included in the model's "
                f"policy configuration. Default intents such as "
                f"'{USER_INTENT_RESTART}' and '{USER_INTENT_BACK}' will "
                f"not trigger actions '{ACTION_RESTART_NAME}' and "
                f"'{ACTION_BACK_NAME}'.",
                docs=DOCS_URL_DEFAULT_ACTIONS,
            )

    def _raise_if_domain_contains_form_names_but_no_rule_policy_given(
        self, domain: Domain
    ) -> None:
        """Validates that there exists a rule policy if forms are defined.

        Raises:
            `InvalidConfigException` if domain and rule policies do not match
        """
        contains_rule_policy = any(
            schema_node
            for schema_node in self._graph_schema.nodes.values()
            if issubclass(schema_node.uses, RulePolicy)
        )

        if domain.form_names and not contains_rule_policy:
            raise InvalidDomain(
                "You have defined a form action, but have not added the "
                f"'{RulePolicy.__name__}' to your policy ensemble. "
                f"Either remove all forms from your domain or add the "
                f"'{RulePolicy.__name__}' to your policy configuration."
            )

    def _raise_if_dm1_and_calm_policies_used_at_the_same_time(self) -> None:
        """Validates that there exists a rule policy if forms are defined.

        Raises:
            `InvalidConfigException` if domain and rule policies do not match
        """
        contains_flow_policy = any(
            schema_node
            for schema_node in self._graph_schema.nodes.values()
            if issubclass(schema_node.uses, FlowPolicy)
        )

        if not contains_flow_policy:
            return

        # if we are using a flow policy, there shouldn't be any dm1 policies
        dm1_policies_list = [
            RulePolicy,
            MemoizationPolicy,
            AugmentedMemoizationPolicy,
        ]
        if TED_POLICY_AVAILABLE:
            dm1_policies_list.append(TEDPolicy)
        if UNEXPECTED_INTENT_POLICY_AVAILABLE:
            dm1_policies_list.append(UnexpecTEDIntentPolicy)
        dm1_policies = tuple(dm1_policies_list)

        for schema_node in self._graph_schema.nodes.values():
            if issubclass(schema_node.uses, dm1_policies):
                raise InvalidConfigException(
                    f"Adding `{schema_node.uses.__name__}` and `FlowPolicy` to "
                    f"the policy ensemble is currently not supported. Please "
                    f"refer to our documentation to understand more about "
                    f"this constraint. https://rasa.com/docs/rasa-pro/concepts/policies"
                )

    def _raise_if_a_rule_policy_is_incompatible_with_domain(
        self, domain: Domain
    ) -> None:
        """Validates the rule policies against the domain.

        Raises:
            `InvalidDomain` if domain and rule policies do not match
        """
        for schema_node in self._graph_schema.nodes.values():
            if schema_node.uses == RulePolicy:
                RulePolicy.raise_if_incompatible_with_domain(
                    config=schema_node.config, domain=domain
                )

    def _validate_policy_priorities(self) -> None:
        """Checks if every policy has a valid priority value.

        A policy must have a priority value. The priority values of
        the policies used in the configuration should be unique.

        Raises:
            `InvalidConfigException` if any of the policies doesn't have a priority
        """
        priority_dict = defaultdict(list)
        for schema_node in self._policy_schema_nodes:
            default_config = schema_node.uses.get_default_config()
            if POLICY_PRIORITY not in default_config:
                raise InvalidConfigException(
                    f"Found a policy {schema_node.uses.__name__} which has no "
                    f"priority. Every policy must have a priority value which you "
                    f"can set in the `get_default_config` method of your policy."
                )
            default_priority = default_config[POLICY_PRIORITY]
            priority = schema_node.config.get(POLICY_PRIORITY, default_priority)
            priority_dict[priority].append(schema_node.uses)

        for k, v in priority_dict.items():
            if len(v) > 1:
                rasa.shared.utils.io.raise_warning(
                    f"Found policies {_types_to_str(v)} with same priority {k} "
                    f"in PolicyEnsemble. When personalizing "
                    f"priorities, be sure to give all policies "
                    f"different priorities.",
                    docs=DOCS_URL_NLU_BASED_POLICIES,
                )

    def _warn_if_rule_based_data_is_unused_or_missing(
        self, story_graph: StoryGraph
    ) -> None:
        """Warns if rule-data is unused or missing.

        Args:
            story_graph: a story graph (core training data)
        """
        consuming_rule_data = any(
            cast(Policy, policy_node.uses).supported_data()
            in [SupportedData.RULE_DATA, SupportedData.ML_AND_RULE_DATA]
            for policy_node in self._policy_schema_nodes
        )

        # Reminder: We generate rule trackers via:
        #   rasa/shared/core/generator/...
        #   .../TrainingDataGenerator/_generate_rule_trackers/
        contains_rule_tracker = any(
            isinstance(step, RuleStep) for step in story_graph.ordered_steps()
        )

        if consuming_rule_data and not contains_rule_tracker:
            rasa.shared.utils.io.raise_warning(
                f"Found a rule-based policy in your configuration but "
                f"no rule-based training data. Please add rule-based "
                f"stories to your training data or "
                f"remove the rule-based policy "
                f"(`{RulePolicy.__name__}`) from your "
                f"your configuration.",
                docs=DOCS_URL_RULES,
            )
        elif not consuming_rule_data and contains_rule_tracker:
            rasa.shared.utils.io.raise_warning(
                f"Found rule-based training data but no policy supporting rule-based "
                f"data. Please add `{RulePolicy.__name__}` "
                f"or another rule-supporting "
                f"policy to the `policies` section in `{DEFAULT_CONFIG_PATH}`.",
                docs=DOCS_URL_RULES,
            )
