import json
import logging
import os
from collections import defaultdict
from types import ModuleType
from typing import Any, Dict, List, Optional, Set, Union

import mlflow
import yaml
from mlflow.models import Model, ModelSignature
from mlflow.utils.file_utils import TempDir, YamlSafeDumper, read_yaml
from pyspark.sql import DataFrame
from pyspark.sql.functions import struct

from feature_store.entities.feature_column_info import FeatureColumnInfo
from feature_store.entities.feature_function import FeatureFunction
from feature_store.entities.feature_lookup import FeatureLookup
from feature_store.entities.feature_spec import FeatureSpec
from feature_store.entities.training_set import TrainingSet
from feature_store.spark_client.spark_client import SparkClient

from feature_store.constants.constants import (
    _NO_RESULT_TYPE_PASSED,
    _PREBUILT_ENV_URI,
    _USE_SPARK_NATIVE_JOIN,
    _WARN,
    MODEL_DATA_PATH_ROOT,
    PREDICTION_COLUMN_NAME,
)

from feature_store.utils import common_utils, training_set_utils
from feature_store.utils.feature_spec_utils import convert_to_yaml_string

_logger = logging.getLogger(__name__)

FEATURE_SPEC_GRAPH_MAX_COLUMN_INFO = 1000


class TrainingSetClient:
    def __init__(
        self,
        spark_client: SparkClient
    ):
        self._spark_client = spark_client

    def create_training_set(
        self,
        feature_spec: FeatureSpec,
        feature_column_infos: List[FeatureColumnInfo],
        label_names: List[str],
        df: DataFrame,
        ft_metadata: training_set_utils._FeatureTableMetadata,
        kwargs,
    ):
        uc_function_infos = training_set_utils.get_uc_function_infos(
            self._spark_client,
            {odci.udf_name for odci in feature_spec.on_demand_column_infos},
        )

        # TODO(divyagupta-db): Move validation from _validate_join_feature_data in feature_lookup_utils.py
        #  to a helper function called here and in score_batch.

        # Add consumer of each feature and instrument as final step
        consumer_feature_table_map = defaultdict(list)
        for feature in feature_column_infos:
            consumer_feature_table_map[feature.table_name].append(feature.feature_name)
        consumed_udf_names = [f.udf_name for f in feature_spec.function_infos]

        # Spark query planning is known to cause spark driver to crash if there are many feature tables to PiT join.
        # See https://docs.google.com/document/d/1EyA4vvlWikTJMeinsLkxmRAVNlXoF1eqoZElOdqlWyY/edit
        # So we disable native join by default.
        training_set_utils.warn_if_non_photon_for_native_spark(
            kwargs.get(_USE_SPARK_NATIVE_JOIN, False), self._spark_client
        )
        return TrainingSet(
            feature_spec,
            df,
            label_names,
            ft_metadata.feature_table_metadata_map,
            ft_metadata.feature_table_data_map,
            uc_function_infos,
            kwargs.get(_USE_SPARK_NATIVE_JOIN, False),
        )

    def create_training_set_from_feature_lookups(
        self,
        df: DataFrame,
        feature_lookups: List[Union[FeatureLookup, FeatureFunction]],
        label: Union[str, List[str], None],
        exclude_columns: List[str],
        **kwargs,
    ) -> TrainingSet:

        # 获取特征查找列表和特征函数列表
        features = feature_lookups
        feature_lookups = [f for f in features if isinstance(f, FeatureLookup)]
        feature_functions = [f for f in features if isinstance(f, FeatureFunction)]

        # 如果未提供标签，则用空列表初始化label_names
        label_names = common_utils.as_list(label, [])
        del label

        # 校验数据集和标签
        training_set_utils.verify_df_and_labels(df, label_names, exclude_columns)

        # 获取特征表元数据
        ft_metadata = training_set_utils.get_table_metadata(
            self._spark_client,
            {fl.table_name for fl in feature_lookups}
        )

        column_infos = training_set_utils.get_column_infos(
            feature_lookups,
            feature_functions,
            ft_metadata,
            df_columns=df.columns,
            label_names=label_names,
        )

        training_set_utils.validate_column_infos(
            self._spark_client,
            ft_metadata,
            column_infos.source_data_column_infos,
            column_infos.feature_column_infos,
            column_infos.on_demand_column_infos,
            label_names,
        )

        # Build feature_spec locally for comparison with the feature spec yaml generated by the
        # FeatureStore backend. This will be removed once the migration is validated.
        feature_spec = training_set_utils.build_feature_spec(
            feature_lookups,
            ft_metadata,
            column_infos,
            exclude_columns
        )

        return self.create_training_set(
            feature_spec,
            column_infos.feature_column_infos,
            label_names,
            df,
            ft_metadata,
            kwargs=kwargs,
        )





    def create_feature_spec(
        self,
        name: str,
        features: List[Union[FeatureLookup, FeatureFunction]],
        sparkClient: SparkClient,
        exclude_columns: List[str] = [],
    ) -> FeatureSpec:

        feature_lookups = [f for f in features if isinstance(f, FeatureLookup)]
        feature_functions = [f for f in features if isinstance(f, FeatureFunction)]

        # Maximum of 100 FeatureFunctions is supported
        if len(feature_functions) > training_set_utils.MAX_FEATURE_FUNCTIONS:
            raise ValueError(
                f"A maximum of {training_set_utils.MAX_FEATURE_FUNCTIONS} FeatureFunctions are supported."
            )

        # Get feature table metadata and column infos
        ft_metadata = training_set_utils.get_table_metadata(
            self._spark_client,
            {fl.table_name for fl in feature_lookups}
        )
        column_infos = training_set_utils.get_column_infos(
            feature_lookups,
            feature_functions,
            ft_metadata,
        )

        column_infos = training_set_utils.add_inferred_source_columns(column_infos)

        training_set_utils.validate_column_infos(
            self._spark_client,
            ft_metadata,
            column_infos.source_data_column_infos,
            column_infos.feature_column_infos,
            column_infos.on_demand_column_infos,
        )

        feature_spec = training_set_utils.build_feature_spec(
            feature_lookups,
            ft_metadata,
            column_infos,
            exclude_columns
        )

        return feature_spec
