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

import mlflow
from mlflow.models import Model
from mlflow.utils.file_utils import TempDir, read_yaml
from pyspark.sql import DataFrame

from wedata.feature_store.constants import constants
from wedata.feature_store.entities.feature_function import FeatureFunction
from wedata.feature_store.entities.feature_lookup import FeatureLookup
from wedata.feature_store.entities.feature_spec import FeatureSpec
from wedata.feature_store.entities.training_set import TrainingSet
from wedata.feature_store.spark_client.spark_client import SparkClient

from wedata.feature_store.constants.constants import (
    _NO_RESULT_TYPE_PASSED,
    _USE_SPARK_NATIVE_JOIN
)

from wedata.feature_store.utils import common_utils, training_set_utils, uc_utils
from wedata.feature_store.utils.signature_utils import get_mlflow_signature_from_feature_spec, \
    drop_signature_inputs_and_invalid_params

_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,
        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},
        )

        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)]

        # 最多支持100个FeatureFunctions
        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."
            )

        # 如果未提供标签，则用空列表初始化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,
            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


    def log_model(
            self,
            model: Any,
            artifact_path: str,
            *,
            flavor: ModuleType,
            training_set: Optional[TrainingSet],
            registered_model_name: Optional[str],
            model_registry_uri: Optional[str],
            await_registration_for: int,
            infer_input_example: bool,
            **kwargs,
    ):
        # Validate only one of the training_set arguments is provided.
        # Retrieve the FeatureSpec, then remove training_set
        if (training_set is None):
            raise ValueError(
                "'training_set' must be provided, but not ."
            )
        # Retrieve the FeatureSpec and then reformat tables in local metastore to 2L before serialization.
        # This will make sure the format of the feature spec with local metastore tables is always consistent.
        # 获取所有表的标签，这里databricks做了标签的规范化
        if training_set:
            print(f'training_set.feature_spec:{training_set.feature_spec}')
            # 判断所有表是否都为dlc表
            all_tables = True
            print(f'all_uc_tables:{all_tables}')
            # training_set.feature_spec is guaranteed to be 3L from FeatureStoreClient.create_training_set.
            feature_spec = uc_utils.get_feature_spec_with_reformat_full_table_names(
                training_set.feature_spec
            )
            print(f'feature_spec:{feature_spec}')

            label_type_map = training_set._label_data_types
            print(f'feature_spec:{feature_spec}')

            labels = training_set._labels

            df_head = training_set._df.head()

        # del training_set

        override_output_schema = kwargs.pop("output_schema", None)
        params = kwargs.pop("params", {})
        params["result_type"] = params.get("result_type", _NO_RESULT_TYPE_PASSED)
        # Signatures will ony be supported for UC-table-only models to
        # mitigate new online scoring behavior from being a breaking regression for older
        # models.
        try:
            # 为所有表的列打标记，输入列为所有特征
            if all_tables:
                signature = get_mlflow_signature_from_feature_spec(
                    feature_spec, label_type_map, override_output_schema, params
                )
            else:
                _logger.warning(
                    "Model could not be logged with a signature because the training set uses feature tables in "
                    "Hive Metastore. Migrate the feature tables to Unity Catalog for model to be logged "
                    "with a signature. "
                )
                signature = None
        except Exception as e:
            _logger.warning(f"Model could not be logged with a signature: {e}")
            signature = None

        with TempDir() as tmp_location:
            # wedata data_path路径,改为记录表路径,遍历表名，生成数组
            data_path = []
            for table_info in training_set.feature_spec.table_infos:
                table_name = common_utils.build_full_table_name(table_info.table_name)
                data_path.append(table_name)

            raw_mlflow_model = Model(
                signature=drop_signature_inputs_and_invalid_params(signature)
            )
            raw_model_path = os.path.join(tmp_location.path(), "constants.RAW_MODEL_FOLDER")

            if flavor.FLAVOR_NAME != mlflow.pyfunc.FLAVOR_NAME:
                flavor.save_model(
                    model, raw_model_path, mlflow_model=raw_mlflow_model, **kwargs
                )
            else:
                flavor.save_model(
                    raw_model_path,
                    mlflow_model=raw_mlflow_model,
                    python_model=model,
                    **kwargs,
                )
            if not "python_function" in raw_mlflow_model.flavors:
                raise ValueError(
                    f"FeatureStoreClient.log_model does not support '{flavor.__name__}' "
                    f"since it does not have a python_function model flavor."
                )

            # Re-use the conda environment from the raw model for the packaged model. Later, we may
            # add an additional requirement for the Feature Store library. At the moment, however,
            # the databricks-feature-store package is not available via conda or pip.
            model_env = raw_mlflow_model.flavors["python_function"][mlflow.pyfunc.ENV]
            if isinstance(model_env, dict):
                # mlflow 2.0 has multiple supported environments
                conda_file = model_env[mlflow.pyfunc.EnvType.CONDA]
            else:
                conda_file = model_env

            conda_env = read_yaml(raw_model_path, conda_file)

            # Check if databricks-feature-lookup version is specified in conda_env
            lookup_client_version_specified = False
            for dependency in conda_env.get("dependencies", []):
                if isinstance(dependency, dict):
                    for pip_dep in dependency.get("pip", []):
                        if pip_dep.startswith(
                                constants.FEATURE_LOOKUP_CLIENT_PIP_PACKAGE
                        ):
                            lookup_client_version_specified = True
                            break

            # If databricks-feature-lookup version is not specified, add default version
            if not lookup_client_version_specified:
                # Get the pip package string for the databricks-feature-lookup client
                default_databricks_feature_lookup_pip_package = common_utils.pip_depependency_pinned_major_version(
                    pip_package_name=constants.FEATURE_LOOKUP_CLIENT_PIP_PACKAGE,
                    major_version=constants.FEATURE_LOOKUP_CLIENT_MAJOR_VERSION,
                )
                common_utils.add_mlflow_pip_depependency(
                    conda_env, default_databricks_feature_lookup_pip_package
                )

            try:
                if df_head is not None and infer_input_example:
                    input_example = df_head.asDict()
                else:
                    input_example = None
            except Exception:
                input_example = None

            # feature_spec.save(data_path)

            # Log the packaged model. If no run is active, this call will create an active run.
            print(f'artifact_path:{artifact_path},data_path:{data_path},conda_env:{conda_env},'
                  f'signature:{signature},input_example:{input_example}');

            mlflow.pyfunc.log_model(
                artifact_path=artifact_path,
                loader_module=constants.MLFLOW_MODEL_NAME,
                data_path=json.dumps(data_path),
                code_path=None,
                conda_env=conda_env,
                signature=signature,
                input_example=input_example,
            )
        if registered_model_name is not None:
            # The call to mlflow.pyfunc.log_model will create an active run, so it is safe to
            # obtain the run_id for the active run.
            run_id = mlflow.tracking.fluent.active_run().info.run_id

            # If the user provided an explicit model_registry_uri when constructing the FeatureStoreClient,
            # we respect this by setting the registry URI prior to reading the model from Model
            # Registry.
        if model_registry_uri is not None:
             # This command will override any previously set registry_uri.
            mlflow.set_registry_uri(model_registry_uri)

            mlflow.register_model(
                "runs:/%s/%s" % (run_id, artifact_path),
                registered_model_name,
                await_registration_for=await_registration_for,
            )
