from typing import Dict
from typing import List
from typing import Tuple

import attrs

from tecton_core import specs
from tecton_core.fco_container import FcoContainer
from tecton_core.feature_definition_wrapper import FeatureDefinitionWrapper
from tecton_core.query_consts import UDF_INTERNAL
from tecton_proto.args import pipeline_pb2


@attrs.frozen
class _Key:
    name: str
    namespace: str


@attrs.frozen
class FeatureDefinitionAndJoinConfig:
    """
    A Feature Definition Wrapper and its associated join configuration.

    :param feature_definition: A FeatureDefinitionWrapper.
    :param join_keys: The mapping from FeatureService's join keys to FeatureView/FeatureTable's join keys.
    :param namespace: The namespace.
    """

    feature_definition: FeatureDefinitionWrapper
    name: str
    # Not a dict to account for multi mapping, though we may not handle multi mapping properly everywhere
    join_keys: List[Tuple[str, str]]
    namespace: str
    features: List[str]

    @classmethod
    def from_feature_definition(cls, feature_definition: FeatureDefinitionWrapper) -> "FeatureDefinitionAndJoinConfig":
        join_keys = [(join_key, join_key) for join_key in feature_definition.join_keys]
        return cls(
            feature_definition=feature_definition,
            name=feature_definition.name,
            join_keys=join_keys,
            namespace=feature_definition.name,
            features=feature_definition.features,
        )

    @classmethod
    def from_feature_set_item_spec(
        cls, feature_set_item: specs.FeatureSetItemSpec, fco_container: FcoContainer
    ) -> "FeatureDefinitionAndJoinConfig":
        """
        :param feature_set_item: A FeatureSetItem spec.
        :param fco_container: Contains all FSI dependencies (transitively), e.g., FV, Entities, DS-es, Transformations
        """

        join_keys = [(i.spine_column_name, i.feature_view_column_name) for i in feature_set_item.join_key_mappings]
        fv_spec = fco_container.get_by_id(feature_set_item.feature_view_id)
        feature_definition = FeatureDefinitionWrapper(fv_spec, fco_container)
        return cls(
            feature_definition=feature_definition,
            name=feature_definition.name,
            join_keys=join_keys,
            namespace=feature_set_item.namespace,
            features=list(feature_set_item.feature_columns),
        )

    def _key(self) -> _Key:
        return _Key(namespace=self.namespace or "", name=self.name)


@attrs.define
class FeatureSetConfig:
    """A wrapper over a list of FeatureDefinitionAndJoinConfigs.

    Used for Feature Service and Feature Definition query construction. Needed for Feature Definition queries because
    some Feature Definitions (namely ODFVs with FV inputs) may require specs from multiple Feature Definitions to
    construct queries.
    """

    definitions_and_configs: List[FeatureDefinitionAndJoinConfig] = attrs.field(factory=list)

    @property  # type: ignore
    def feature_definitions(self) -> List[FeatureDefinitionWrapper]:
        """
        Returns the FeatureViews/FeatureTables enclosed in this FeatureSetConfig.
        """
        return [config.feature_definition for config in self.definitions_and_configs]

    @property
    def features(self) -> List[str]:
        """
        Returns the features generated by the enclosed feature definitions.
        """
        return [
            features
            for config in self.definitions_and_configs
            for features in FeatureSetConfig._get_full_feature_names(config)
        ]

    @property
    def join_keys(self) -> List[str]:
        """
        Returns the join keys used across all feature definitions.
        """
        keys = set()
        for dac in self.definitions_and_configs:
            for jk_pair in dac.join_keys:
                service_key = jk_pair[0]
                if service_key not in keys:
                    keys.add(service_key)
        return list(keys)

    @property
    def request_context_keys(self) -> List[str]:
        all_keys = sum([fd.request_context_keys for fd in self.feature_definitions], [])
        unique_keys = list(set(all_keys))
        return unique_keys

    @staticmethod
    def _get_full_feature_names(config: FeatureDefinitionAndJoinConfig) -> List[str]:
        return [
            config.namespace + config.feature_definition.namespace_separator + feature if config.namespace else feature
            for feature in config.features
        ]

    @classmethod
    def from_feature_definition(cls, feature_definition: FeatureDefinitionWrapper) -> "FeatureSetConfig":
        definitions_and_configs = [FeatureDefinitionAndJoinConfig.from_feature_definition(feature_definition)]

        if feature_definition.is_on_demand:
            inputs = find_dependent_feature_set_items(
                feature_definition.fco_container,
                feature_definition.pipeline.root,
                visited_inputs={},
                fv_id=feature_definition.id,
            )
            definitions_and_configs.extend(inputs)

        return FeatureSetConfig(definitions_and_configs=definitions_and_configs)

    @classmethod
    def from_feature_service_spec(
        cls, feature_service_spec: specs.FeatureServiceSpec, fco_container: FcoContainer
    ) -> "FeatureSetConfig":
        """
        :param feature_service_spec: A feature service spec.
        :param fco_container: Contains all FSI dependencies (transitively), e.g., FVs, Entities, DS-es, Transformations
        """

        definitions_and_configs = [
            FeatureDefinitionAndJoinConfig.from_feature_set_item_spec(feature_set_item, fco_container)
            for feature_set_item in feature_service_spec.feature_set_items
        ]

        # Add dependent feature views into the FeatureSetConfig, uniquely per odfv
        # The namespaces of the dependencies have _udf_internal in the name and are filtered out before
        # being returned by TectonContext.execute()
        visited_feature_view_ids = set()
        for feature_set_item in feature_service_spec.feature_set_items:
            if feature_set_item.feature_view_id in visited_feature_view_ids:
                continue
            visited_feature_view_ids.add(feature_set_item.feature_view_id)

            fv_spec = fco_container.get_by_id(feature_set_item.feature_view_id)
            if not isinstance(fv_spec, specs.FeatureViewSpec):
                TypeError(f"Expected a feature view type. {fv_spec}")

            if isinstance(fv_spec, specs.OnDemandFeatureViewSpec):
                dependent_items = find_dependent_feature_set_items(
                    fco_container,
                    fv_spec.pipeline.root,
                    visited_inputs={},
                    fv_id=fv_spec.id,
                )
                definitions_and_configs.extend(dependent_items)

        return cls(definitions_and_configs=definitions_and_configs)


def find_dependent_feature_set_items(
    fco_container: FcoContainer, node: pipeline_pb2.PipelineNode, visited_inputs: Dict[str, bool], fv_id: str
) -> List[FeatureDefinitionAndJoinConfig]:
    if node.HasField("feature_view_node"):
        if node.feature_view_node.input_name in visited_inputs:
            return []
        visited_inputs[node.feature_view_node.input_name] = True

        fv_spec = fco_container.get_by_id_proto(node.feature_view_node.feature_view_id)
        fd = FeatureDefinitionWrapper(fv_spec, fco_container)

        join_keys = []
        overrides = {
            colpair.feature_column: colpair.spine_column
            for colpair in node.feature_view_node.feature_view.override_join_keys
        }
        for join_key in fd.join_keys:
            potentially_overriden_key = overrides.get(join_key, join_key)
            join_keys.append((potentially_overriden_key, join_key))

        cfg = FeatureDefinitionAndJoinConfig(
            feature_definition=fd,
            name=fd.name,
            join_keys=join_keys,
            namespace=f"{UDF_INTERNAL}_{node.feature_view_node.input_name}_{fv_id}",
            features=list(node.feature_view_node.feature_view.features) or fd.features,
        )

        return [cfg]
    elif node.HasField("transformation_node"):
        ret: List[FeatureDefinitionAndJoinConfig] = []
        for child in node.transformation_node.inputs:
            ret = ret + find_dependent_feature_set_items(fco_container, child.node, visited_inputs, fv_id)
        return ret
    return []
