import os
from typing import Any, Dict, List, Type, Union

import mlflow
from google.protobuf.json_format import MessageToDict, ParseDict
from mlflow.utils.file_utils import TempDir
from typing_extensions import override

from databricks.ml_features_common.entities._feature_spec_base import _FeatureSpecBase
from databricks.ml_features_common.entities.column_info import ColumnInfo
from databricks.ml_features_common.entities.data_type import DataType
from databricks.ml_features_common.entities.feature_column_info import FeatureColumnInfo
from databricks.ml_features_common.entities.feature_spec_constants import (
    BOUND_TO,
    DATA_TYPE,
    FEATURE_COLUMN_INFO,
    FEATURE_STORE,
    INCLUDE,
    INPUT_BINDINGS,
    INPUT_COLUMNS,
    INPUT_FUNCTIONS,
    INPUT_TABLES,
    NAME,
    ON_DEMAND_COLUMN_INFO,
    ON_DEMAND_FEATURE,
    OUTPUT_NAME,
    PARAMETER,
    SERIALIZATION_VERSION,
    SOURCE,
    SOURCE_DATA_COLUMN_INFO,
    TABLE_NAME,
    TOPOLOGICAL_ORDERING,
    TRAINING_DATA,
    UDF_NAME,
)
from databricks.ml_features_common.entities.feature_table_info import FeatureTableInfo
from databricks.ml_features_common.entities.feature_tables_for_serving import (
    FeatureTablesForServing,
)
from databricks.ml_features_common.entities.function_info import FunctionInfo
from databricks.ml_features_common.entities.on_demand_column_info import (
    OnDemandColumnInfo,
)
from databricks.ml_features_common.entities.source_data_column_info import (
    SourceDataColumnInfo,
)
from databricks.ml_features_common.protos.feature_spec_pb2 import (
    FeatureSpec as ProtoFeatureSpec,
)
from databricks.ml_features_common.utils import utils_common
from databricks.ml_features_common.utils.utils_common import is_artifact_uri
from databricks.ml_features_common.utils.yaml_utils import read_yaml, write_yaml

# Change log for serialization version. Please update for each serialization version.
# 1. Initial.
# 2. (2021/06/16): Record feature_store_client_version to help us make backward compatible changes in the future.
# 3. (2021/08/25): Record table_id to handle feature table lineage stability if tables are deleted.
# 4. (2021/09/25): Record timestamp_lookup_key to handle point-in-time lookups.
# 5. (2021/02/15): Record include flag for column info if False.
#                  Record input functions as FunctionInfo and function computation as OnDemandColumnInfo.
#                  Remove redundant fields: table_name from table_infos, output_name from column_infos.
# 6. (2023/04/21): Record lookback_window in table info for point-in-time lookups.
# 7. (2023/05/05): Record the Spark data type for all columns to track model signatures.
# 8. (2023/08/14): Record the topological_ordering for all columns to support chained transform and lookup.
# 9. (2023/09/11): Change the type of lookback_window from int to double for sub-second values


class FeatureSpec(_FeatureSpecBase):
    def __init__(
        self,
        column_infos: List[ColumnInfo],
        table_infos: List[FeatureTableInfo],
        function_infos: List[FunctionInfo],
        workspace_id: int,
        feature_store_client_version: str,
        serialization_version: int,
    ):
        super().__init__(
            serialization_version, workspace_id, feature_store_client_version
        )
        self._column_infos = column_infos
        self._table_infos = table_infos
        self._function_infos = function_infos

        # Perform validations
        self._validate_column_infos()
        self._validate_table_infos()
        self._validate_function_infos()

    def _validate_column_infos(self):
        if not self.column_infos:
            raise ValueError("column_infos must be non-empty.")

        for column_info in self.column_infos:
            if not isinstance(column_info, ColumnInfo):
                raise ValueError(
                    f"Expected all elements of column_infos to be instances of ColumnInfo. "
                    f"'{column_info}' is of the wrong type."
                )
            if (
                self._serialization_version >= 8
                and column_info.topological_ordering is not None
            ):
                ordering = column_info.topological_ordering
                if not isinstance(ordering, int) or ordering < 0:
                    raise ValueError(
                        "The topological_ordering of column_info must be non non-negative integers."
                    )

    def _validate_table_infos(self):
        if self.table_infos is None:
            raise ValueError("Internal Error: table_infos must be provided.")

        # table_infos should not be duplicated
        utils_common.validate_strings_unique(
            [table_info.table_name for table_info in self.table_infos],
            "Internal Error: Expect all table_names in table_infos to be unique. Found duplicates {}",
        )

        # Starting FeatureSpec v3, unique table names in table_infos must match those in column_infos.
        if self.serialization_version >= 3:
            unique_table_names = set(
                [table_info.table_name for table_info in self.table_infos]
            )
            unique_column_table_names = set(
                [fci.table_name for fci in self.feature_column_infos]
            )
            if unique_table_names != unique_column_table_names:
                raise Exception(
                    f"Internal Error: table_names from table_infos {sorted(unique_table_names)} "
                    f"must match those from column_infos {sorted(unique_column_table_names)}"
                )

    def _validate_function_infos(self):
        if self.function_infos is None:
            raise ValueError("Internal Error: function_infos must be provided.")

        # function_infos should not be duplicated
        utils_common.validate_strings_unique(
            [function_info.udf_name for function_info in self.function_infos],
            "Internal Error: Expect all udf_names in function_infos to be unique. Found duplicates {}",
        )

        # Unique UDF names in function_infos must match those in column_infos.
        # No version check is required as both fields were added simultaneously in FeatureSpec v5.
        unique_udf_names = set(
            [function_info.udf_name for function_info in self.function_infos]
        )
        unique_column_udf_names = set(
            [odci.udf_name for odci in self.on_demand_column_infos]
        )
        if unique_udf_names != unique_column_udf_names:
            raise Exception(
                f"Internal Error: udf_names from function_infos {sorted(unique_udf_names)} "
                f"must match those from column_infos {sorted(unique_column_udf_names)}"
            )

    @property
    def column_infos(self):
        return self._column_infos

    @property
    def table_infos(self):
        return self._table_infos

    @property
    def function_infos(self):
        return self._function_infos

    def _get_infos_of_type(
        self,
        info_type: Union[
            Type[SourceDataColumnInfo],
            Type[FeatureColumnInfo],
            Type[OnDemandColumnInfo],
        ],
    ):
        """
        Helper method to return the ColumnInfo.info subinfo field based on its type.
        """
        return [
            column_info.info
            for column_info in self.column_infos
            if isinstance(column_info.info, info_type)
        ]

    @property
    def source_data_column_infos(self) -> List[SourceDataColumnInfo]:
        return self._get_infos_of_type(SourceDataColumnInfo)

    @property
    def feature_column_infos(self) -> List[FeatureColumnInfo]:
        return self._get_infos_of_type(FeatureColumnInfo)

    @property
    def on_demand_column_infos(self) -> List[OnDemandColumnInfo]:
        return self._get_infos_of_type(OnDemandColumnInfo)

    @classmethod
    def from_proto(cls, feature_spec_proto):
        # Serialization version is not deserialized from the proto as there is currently only one
        # possible version.
        column_infos = [
            ColumnInfo.from_proto(column_info_proto)
            for column_info_proto in feature_spec_proto.input_columns
        ]
        table_infos = [
            FeatureTableInfo.from_proto(table_info_proto)
            for table_info_proto in feature_spec_proto.input_tables
        ]
        function_infos = [
            FunctionInfo.from_proto(function_info_proto)
            for function_info_proto in feature_spec_proto.input_functions
        ]
        return cls(
            column_infos=column_infos,
            table_infos=table_infos,
            function_infos=function_infos,
            workspace_id=feature_spec_proto.workspace_id,
            feature_store_client_version=feature_spec_proto.feature_store_client_version,
            serialization_version=feature_spec_proto.serialization_version,
        )

    def has_feature_renaming(self):
        for fc in self.feature_column_infos:
            if fc.output_name != fc.feature_name:
                return True
        for oc in self.on_demand_column_infos:
            if oc.output_name != oc.udf_name:
                return True
        return False

    def has_default_value(self):
        for fc in self.feature_column_infos:
            if fc.default_value_str:
                return True
        return False

    def has_dag(self):
        """
        Check if the FeatureSpec has a DAG. If a DAG is detected the feature lookup needs
        to be performed in multiple steps.
        """
        source_column_names = set(info.name for info in self.source_data_column_infos)
        for column_info in self.column_infos:
            if isinstance(column_info.info, FeatureColumnInfo):
                for lookup_key in column_info.info.lookup_key:
                    if lookup_key not in source_column_names:
                        return True
            elif isinstance(column_info.info, OnDemandColumnInfo):
                for parameter in column_info.info.input_bindings.keys():
                    if parameter not in source_column_names:
                        return True
        return False

    def to_proto(self):
        proto_feature_spec = ProtoFeatureSpec()
        for column_info in self.column_infos:
            proto_feature_spec.input_columns.append(column_info.to_proto())
        for table_info in self.table_infos:
            proto_feature_spec.input_tables.append(table_info.to_proto())
        for function_info in self.function_infos:
            proto_feature_spec.input_functions.append(function_info.to_proto())
        proto_feature_spec.serialization_version = self.serialization_version
        proto_feature_spec.workspace_id = self.workspace_id
        proto_feature_spec.feature_store_client_version = (
            self._feature_store_client_version
        )
        return proto_feature_spec

    @staticmethod
    def _input_columns_proto_to_yaml_dict(column_info: Dict[str, Any]):
        """
        Converts a single ColumnInfo's proto dict to the expected element in FeatureSpec YAML's input_columns.
        To keep the YAML clean, unnecessary fields are removed (e.g. SourceDataColumnInfo.name field, ColumnInfo.include when True).

        Example of a column_info transformation. Note that "name" and "include" attributes were excluded.
        {"source_data_column_info": {"name": "source_column"}, "include": True} -> {"source_column": {"source": "training_data"}}

        Order of elements in the YAML dict should be:
        1. Attributes present in ColumnInfo.info, using the proto field order
        2. Remaining attributes of ColumnInfo, using the proto field order
        3. Feature Store source type
        """
        # Parse oneof field ColumnInfo.info level attributes as column_info_attributes; record column_name, source
        if SOURCE_DATA_COLUMN_INFO in column_info:
            column_info_attributes = column_info[SOURCE_DATA_COLUMN_INFO]
            # pop NAME attribute and use as the YAML key for this column_info to avoid redundancy in YAML
            column_name, source = column_info_attributes.pop(NAME), TRAINING_DATA
        elif FEATURE_COLUMN_INFO in column_info:
            column_info_attributes = column_info[FEATURE_COLUMN_INFO]
            # pop OUTPUT_NAME attribute and use as the YAML key for this column_info to avoid redundancy in YAML
            column_name, source = column_info_attributes.pop(OUTPUT_NAME), FEATURE_STORE
        elif ON_DEMAND_COLUMN_INFO in column_info:
            column_info_attributes = column_info[ON_DEMAND_COLUMN_INFO]
            # Map InputBindings message dictionary to {parameter: bound_to} KV dictionary if defined
            if INPUT_BINDINGS in column_info_attributes:
                column_info_attributes[INPUT_BINDINGS] = {
                    ib[PARAMETER]: ib[BOUND_TO]
                    for ib in column_info_attributes[INPUT_BINDINGS]
                }
            # pop OUTPUT_NAME attribute and use as the YAML key for this column_info to avoid redundancy in YAML
            column_name, source = (
                column_info_attributes.pop(OUTPUT_NAME),
                ON_DEMAND_FEATURE,
            )
        else:
            raise ValueError(
                f"Expected column_info to be keyed by a valid ColumnInfo.info type. "
                f"'{column_info}' has key '{list(column_info)[0]}'."
            )

        # Parse and insert ColumnInfo level attributes
        # Note: the ordering of fields in the result yaml file is undefined but in reality, they are
        # in the same order as they are added in the column_info_attributes dict.

        # DATA_TYPE is supported starting FeatureSpec v7 and is not guaranteed to exist.
        if DATA_TYPE in column_info:
            column_info_attributes[DATA_TYPE] = column_info[DATA_TYPE]
        if not column_info[INCLUDE]:
            column_info_attributes[INCLUDE] = False
        # TOPOLOGICAL_ORDERING is supported starting FeatureSpec v8.
        if TOPOLOGICAL_ORDERING in column_info:
            column_info_attributes[TOPOLOGICAL_ORDERING] = column_info[
                TOPOLOGICAL_ORDERING
            ]

        # Insert source; return YAML keyed by column_name
        column_info_attributes[SOURCE] = source
        return {column_name: column_info_attributes}

    @override
    def _to_dict(self):
        """
        Convert FeatureSpec to a writeable YAML artifact. Uses MessageToDict to convert FeatureSpec proto to dict.
        Sanitizes and modifies the dict as follows:
        1. Remove redundant or unnecessary information for cleanliness in the YAML
        2. Modifies the dict to be of the format {column_name: column_attributes_dict}

        :return: Sanitized FeatureSpec dictionary of {column_name: column_attributes}
        """
        yaml_dict = MessageToDict(self.to_proto(), preserving_proto_field_name=True)
        yaml_dict[INPUT_COLUMNS] = [
            self._input_columns_proto_to_yaml_dict(column_info)
            for column_info in yaml_dict[INPUT_COLUMNS]
        ]

        if INPUT_TABLES in yaml_dict:
            # pop TABLE_NAME attribute and use as the YAML key for each table_info to avoid redundancy in YAML
            yaml_dict[INPUT_TABLES] = [
                {table_info.pop(TABLE_NAME): table_info}
                for table_info in yaml_dict[INPUT_TABLES]
            ]
        if INPUT_FUNCTIONS in yaml_dict:
            # pop UDF_NAME attribute and use as the YAML key for each table_info to avoid redundancy in YAML
            yaml_dict[INPUT_FUNCTIONS] = [
                {function_info.pop(UDF_NAME): function_info}
                for function_info in yaml_dict[INPUT_FUNCTIONS]
            ]

        # For readability, place SERIALIZATION_VERSION last in the dictionary.
        yaml_dict[SERIALIZATION_VERSION] = yaml_dict.pop(SERIALIZATION_VERSION)
        return yaml_dict

    @staticmethod
    def _input_columns_yaml_to_proto_dict(column_info: Dict[str, Any]):
        """
        Convert the FeatureSpec YAML dictionary to the expected ColumnInfo proto dictionary.

        Example of a column_info transformation.
        {"source_column": {"source": "training_data"}} -> {"source_data_column_info": {"name": "source_column"}}
        """
        if len(column_info) != 1:
            raise ValueError(
                f"Expected column_info dictionary to only have one key, value pair. "
                f"'{column_info}' has length {len(column_info)}."
            )
        column_name, column_data = list(column_info.items())[0]
        if not column_data:
            raise ValueError(
                f"Expected values of '{column_name}' dictionary to be non-empty."
            )
        if SOURCE not in column_data:
            raise ValueError(
                f"Expected values of column_info dictionary to include the source. No source found "
                f"for '{column_name}'."
            )

        # Parse oneof field ColumnInfo.info level attributes
        source = column_data.pop(SOURCE)
        if source == TRAINING_DATA:
            column_data[NAME] = column_name
            column_info_dict = {SOURCE_DATA_COLUMN_INFO: column_data}
        elif source == FEATURE_STORE:
            column_data[OUTPUT_NAME] = column_name
            column_info_dict = {FEATURE_COLUMN_INFO: column_data}
        elif source == ON_DEMAND_FEATURE:
            column_data[OUTPUT_NAME] = column_name
            # Map {parameter_val: bound_to_val} dictionary to InputBindings(parameter, bound_to) message dictionary.
            column_data[INPUT_BINDINGS] = [
                {PARAMETER: parameter, BOUND_TO: bound_to}
                for parameter, bound_to in column_data.get(INPUT_BINDINGS, {}).items()
            ]
            column_info_dict = {ON_DEMAND_COLUMN_INFO: column_data}
        else:
            raise ValueError(
                f"Internal Error: Expected column_info to have source matching oneof ColumnInfo.info. "
                f"'{column_info}' has source of '{source}'."
            )

        # Parse ColumnInfo level attributes
        # TOPOLOGICAL_ORDERING is supported starting FeatureSpec v8.
        if TOPOLOGICAL_ORDERING in column_data:
            column_info_dict[TOPOLOGICAL_ORDERING] = column_data.pop(
                TOPOLOGICAL_ORDERING
            )
        # DATA_TYPE is supported starting FeatureSpec v7 and is not guaranteed to exist.
        if DATA_TYPE in column_data:
            column_info_dict[DATA_TYPE] = column_data.pop(DATA_TYPE)
        # INCLUDE is supported starting FeatureSpec v5 and only present in the YAML when INCLUDE = False
        if INCLUDE in column_data:
            column_info_dict[INCLUDE] = column_data.pop(INCLUDE)
        return column_info_dict

    @classmethod
    def _from_dict(cls, spec_dict):
        """
        Convert YAML artifact to FeatureSpec. Transforms YAML artifact to dict keyed by
        source_data_column_info or feature_column_info, such that ParseDict can convert the dict to
        a proto message, and from_proto can convert the proto message to a FeatureSpec object
        :return: :py:class:`~databricks.ml_features_common.entities.feature_spec.FeatureSpec`
        """
        if INPUT_COLUMNS not in spec_dict:
            raise ValueError(
                f"{INPUT_COLUMNS} must be a key in {cls.FEATURE_ARTIFACT_FILE}."
            )
        if not spec_dict[INPUT_COLUMNS]:
            raise ValueError(
                f"{INPUT_COLUMNS} in {cls.FEATURE_ARTIFACT_FILE} must be non-empty."
            )
        spec_dict[INPUT_COLUMNS] = [
            cls._input_columns_yaml_to_proto_dict(column_info)
            for column_info in spec_dict[INPUT_COLUMNS]
        ]

        # feature_spec.yaml doesn't include input_tables, input_functions if any are true:
        # 1. The YAML is written by an older client that does not support the functionality.
        # 2. The FeatureSpec does not contain FeatureLookups (input_tables), FeatureFunctions (input_functions).
        input_tables = []
        for input_table in spec_dict.get(INPUT_TABLES, []):
            table_name, attributes = list(input_table.items())[0]
            input_tables.append({TABLE_NAME: table_name, **attributes})
        spec_dict[INPUT_TABLES] = input_tables

        input_functions = []
        for input_function in spec_dict.get(INPUT_FUNCTIONS, []):
            udf_name, attributes = list(input_function.items())[0]
            input_functions.append({UDF_NAME: udf_name, **attributes})
        spec_dict[INPUT_FUNCTIONS] = input_functions

        return cls.from_proto(
            ParseDict(spec_dict, ProtoFeatureSpec(), ignore_unknown_fields=True)
        )

    @classmethod
    def _read_file(cls, path: str):
        """
        Read the YAML artifact from a file path.
        """
        parent_dir, file = os.path.split(path)
        spec_dict = read_yaml(parent_dir, file)
        return cls._from_dict(spec_dict)

    @classmethod
    def load(cls, path: str):
        """
        Deprecated: Use feature_spec_utils::load_feature_spec instead.

        Load the FeatureSpec YAML artifact in the provided root directory (at path/feature_spec.yaml).

        :param path: Root path to the YAML artifact. This can be a MLflow artifact path or file path.
        :return: :py:class:`~databricks.ml_features_common.entities.feature_spec.FeatureSpec`
        """
        # Create the full file path to the FeatureSpec.
        path = os.path.join(path, cls.FEATURE_ARTIFACT_FILE)

        if is_artifact_uri(path):
            with TempDir() as tmp_location:
                # Returns a file and not directory since the artifact_uri is a single file.
                local_path = mlflow.artifacts.download_artifacts(
                    artifact_uri=path, dst_path=tmp_location.path()
                )
                return FeatureSpec._read_file(local_path)
        else:
            return FeatureSpec._read_file(path)

    @classmethod
    def from_feature_tables_for_serving(
        cls,
        feature_tables_for_serving: FeatureTablesForServing,
    ) -> "FeatureSpec":
        """
        Constructs a FeatureSpec object from FeatureTablesForServing for declarative features.

        This method extracts feature metadata from online feature tables and constructs a
        FeatureSpec programmatically. It creates ColumnInfo objects for each feature in each
        online feature table, mapping primary keys and timestamp keys, and creates
        FeatureTableInfo objects for each unique table name.

        Notes:
        1. For declarative features, we use the online table name as the table name in the
           FeatureSpec. This is because declarative features might not be materialized in an
           offline table (e.g., for streaming cases), and the table name is only used to look
           up online tables.
        2. We do not support renaming for declarative features. Therefore,
           we assume the declarative feature name is unique and matches the column name in the
           online table. It also matches the column name in the dataframe resulting from the lookup.
           (i.e., feature_name == output_name).
        3. We do not support default values for declarative features.

        :param feature_tables_for_serving: FeatureTablesForServing object containing online
            feature table definitions with features, primary keys, and timestamp keys.
        :param workspace_id: Workspace ID for the FeatureSpec. Defaults to 0.
        :param feature_store_client_version: Feature Store client version. Defaults to empty string.
        :return: FeatureSpec object with column_infos, table_infos, and empty function_infos
            (since declarative features don't use feature functions).
        """

        column_infos = [
            ColumnInfo(
                info=FeatureColumnInfo(
                    table_name=online_feature_table.online_feature_table_name,
                    feature_name=feature.name,
                    lookup_key=[key.name for key in online_feature_table.primary_keys],
                    output_name=feature.name,
                    timestamp_lookup_key=[
                        key.name for key in online_feature_table.timestamp_keys
                    ],
                ),
                include=True,
                data_type=DataType.to_string(feature.data_type)
                if feature.data_type
                else None,
            )
            for online_feature_table in feature_tables_for_serving.online_feature_tables
            for feature in online_feature_table.features
        ]

        # Create FeatureTableInfo objects for each unique table name
        # FeatureSpec validation requires table_infos to match table names in column_infos
        unique_table_names = set(
            online_feature_table.online_feature_table_name
            for online_feature_table in feature_tables_for_serving.online_feature_tables
        )
        table_infos = [
            FeatureTableInfo(
                table_name=table_name,
                table_id=None,
            )
            for table_name in unique_table_names
        ]

        return cls(
            column_infos=column_infos,
            table_infos=table_infos,
            function_infos=[],
            workspace_id=None,
            feature_store_client_version=None,
            serialization_version=cls.SERIALIZATION_VERSION_NUMBER,
        )
