from typing import List
from typing import Optional
from typing import Tuple
from typing import Union

import attr
import pandas
import pyspark

from tecton._internals.sdk_decorators import sdk_public_method
from tecton.interactive.data_frame import DataFrame
from tecton_proto.data.feature_service_pb2 import FeatureSetItem
from tecton_spark.feature_definition_wrapper import FeatureDefinitionWrapper as FeatureDefinition
from tecton_spark.logger import get_logger

logger = get_logger("FeatureSetConfig")


@attr.s(frozen=True, auto_attribs=True)
class _Key(object):
    name: str
    namespace: str


@attr.s(frozen=True, auto_attribs=True)
class FeatureDefinitionAndJoinConfig(object):
    """
    A feature definition and its associated join configuration.

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

    feature_definition: FeatureDefinition
    name: str
    join_keys: List[Tuple[str, str]]
    namespace: str
    features: List[str]

    @classmethod
    def _from_proto(cls, feature_set_item: FeatureSetItem):
        join_keys = [(i.spine_column_name, i.package_column_name) for i in feature_set_item.join_configuration_items]
        feature_definition = FeatureDefinition(feature_set_item.enrichments.feature_view)
        return FeatureDefinitionAndJoinConfig(
            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)

    def to_proto(self) -> FeatureSetItem:
        feature_set_item = FeatureSetItem()
        feature_set_item.namespace = self.namespace
        feature_set_item.feature_columns.extend(self.features)
        feature_set_item.enrichments.feature_view.CopyFrom(self.feature_definition.feature_view_proto)
        for spine_column_name, package_column_name in self.join_keys:
            join_configuration_item = feature_set_item.join_configuration_items.add()
            join_configuration_item.spine_column_name = spine_column_name
            join_configuration_item.package_column_name = package_column_name
        return feature_set_item


class FeatureSetConfig(object):
    """
    Config used to create a :class:`FeatureService`.
    """

    # This stores all the features + all the odfv dependent features
    _definitions_and_configs: List[FeatureDefinitionAndJoinConfig]
    # This stores all the odfv dependent features
    _ondemand_input_definitions_and_configs: List[FeatureDefinitionAndJoinConfig]

    def __init__(self):
        """
        Initialize a new FeatureSetConfig.
        """

        self._definitions_and_configs = []
        self._ondemand_input_definitions_and_configs = []

    def _add(
        self,
        fd: FeatureDefinition,
        namespace: Optional[str] = None,
    ):
        if namespace is None:
            namespace = fd.name

        fd_and_config = FeatureSetConfig._make_definition_and_config(fd, namespace=namespace)
        self._definitions_and_configs.append(fd_and_config)

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

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

    @staticmethod
    def _get_full_feature_names(config: FeatureDefinitionAndJoinConfig):
        return [config.namespace + "." + feature if config.namespace else feature for feature in config.features]

    @staticmethod
    def _make_definition_and_config(fd: FeatureDefinition, namespace: str) -> FeatureDefinitionAndJoinConfig:
        join_keys = [(join_key, join_key) for join_key in fd.join_keys]
        features = fd.features

        return FeatureDefinitionAndJoinConfig(
            feature_definition=fd, name=fd.name, join_keys=join_keys, namespace=namespace, features=features
        )

    def _get_feature_definitions_and_join_configs(self) -> List[FeatureDefinitionAndJoinConfig]:
        """
        Returns a list of all feature definitions and their configs in the insertion order.
        """
        return self._definitions_and_configs

    def _get_ondemand_feature_definitions_and_join_configs(self) -> List[FeatureDefinitionAndJoinConfig]:
        """
        Returns a list of all the odfv dependent feature definitions and their configs in the insertion order.
        """
        return self._ondemand_input_definitions_and_configs

    @sdk_public_method
    def get_feature_dataframe(
        self,
        spine: Union[str, dict, "pyspark.sql.DataFrame", list, "pandas.DataFrame", "pyspark.RDD", None],
        timestamp_key: str,
        include_feature_view_timestamp_columns=False,
        use_materialized_data: bool = True,
    ) -> DataFrame:
        """
        Creates a feature vector dataframe for the specified entities, including the entities' features.

        :param spine: SQL string that fetches the entity ids for which the features will be generated or a dataframe.
        :param timestamp_key: Name of the time column in spine. The column must be of type Spark timestamp.
        :param include_feature_view_timestamp_columns: (Optional) Include timestamp columns for every individual feature definition.
        :param use_materialized_data: (Optional) Use materialized data if materialization is enabled.
        :return: A Tecton DataFrame.
        """
        from tecton.tecton_context import TectonContext

        tc = TectonContext.get_instance()
        return tc.execute(
            spine,
            feature_set_config=self,
            timestamp_key=timestamp_key,
            include_feature_view_timestamp_columns=include_feature_view_timestamp_columns,
            use_materialized_data=use_materialized_data,
        )

    @classmethod
    def _from_protos(cls, feature_set_items):
        config = FeatureSetConfig()
        config._definitions_and_configs = [FeatureDefinitionAndJoinConfig._from_proto(p) for p in feature_set_items]
        return config
