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

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

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.feature_table import FeatureTable
from wedata.feature_store.entities.training_set import TrainingSet
from wedata.feature_store.mlflow_model import _FeatureStoreModelWrapper
from wedata.feature_store.spark_client.spark_client import SparkClient
from wedata.feature_store.utils import validation_utils
from wedata.feature_store.entities.feature_table import FeatureTable

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

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,
    ):
        # 验证training_set参数是否提供
        if (training_set is None):
            raise ValueError(
                "'training_set' must be provided, but not ."
            )

        # 获取特征规格并重新格式化表名
        # training_set.feature_spec保证来自FeatureStoreClient.create_training_set的3L格式
        feature_spec = uc_utils.get_feature_spec_with_reformat_full_table_names(
            training_set.feature_spec
        )

        # 获取标签类型映射和标签
        label_type_map = training_set._label_data_types

        # 收集所有特征列名
        feature_columns = [
            feature_column.output_name
            for feature_column in feature_spec.feature_column_infos
        ]
        df_head = training_set.load_df().select(*feature_columns).head()

        # 处理输出模式和参数
        override_output_schema = kwargs.pop("output_schema", None)
        params = kwargs.pop("params", {})
        params["result_type"] = params.get("result_type", _NO_RESULT_TYPE_PASSED)

        # 尝试获取MLflow签名
        try:
            signature = get_mlflow_signature_from_feature_spec(
                feature_spec, label_type_map, override_output_schema, params
            )
        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 = os.path.join(tmp_location.path(), "feature_store")
            os.makedirs(data_path, exist_ok=True)

            # 创建原始MLflow模型
            raw_mlflow_model = Model(
                signature=drop_signature_inputs_and_invalid_params(signature)
            )
            raw_model_path = os.path.join(data_path, constants.RAW_MODEL_FOLDER)

            # 根据flavor类型保存模型
            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,
                )

            # 验证模型是否支持python_function flavor
            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."
                )

            # 获取并处理conda环境配置
            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)
            #TODO 暂时不需要databricks-feature-lookup这个包，会导致 python 环境创建失败
            # # 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
            #TODO 暂时不需要databricks-feature-lookup这个包，会导致 python 环境创建失败
            # # 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
            #     )

            # 尝试创建输入示例
            input_example = None
            try:
                if df_head is not None and infer_input_example:
                    input_example = df_head.asDict()
            except Exception:
                pass

            feature_spec.save(data_path)

            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,
                python_model=_FeatureStoreModelWrapper(model),
                # data_path=data_path,
                artifacts={"feature_store": data_path},
                code_path=None,
                conda_env=conda_env,
                signature=signature,
                input_example=input_example,
                registered_model_name=registered_model_name
            )

            # mlflow.pyfunc.log_model(
            #     artifact_path=artifact_path,
            #     loader_module=constants.MLFLOW_MODEL_NAME,
            #     data_path=data_path,
            #     conda_env=conda_env,
            #     signature=signature,
            #     input_example=input_example,
            # )

        # 注册模型
        # if registered_model_name is not None:
        #     run_id = mlflow.tracking.fluent.active_run().info.run_id
        # if model_registry_uri is not None:
        #     mlflow.set_registry_uri(model_registry_uri)
        #
        # mlflow.register_model(
        #     f"runs:/{run_id}/{artifact_path}",
        #     registered_model_name,
        #     await_registration_for=await_registration_for,
        # )
        #
        # print(f"Model registered successfully: {registered_model_name}")

        # # 验证模型是否已注册
        # from mlflow.tracking import MlflowClient
        # client = MlflowClient()
        # model_version = client.get_latest_versions(registered_model_name, stages=["None"])[0]
        # print(f"Registered model version: {model_version.version}")

    def score_batch(
        self,
        model_uri: Optional[str],
        df: DataFrame,
        result_type: str,
        env_manager: Optional[str] = None,
        local_uri: Optional[str] = None,
        params: Optional[dict[str, Any]] = None,
        **kwargs,
    ) -> DataFrame:
        # TODO:ML 待确定是否需要
        # req_context = RequestContext(request_context.SCORE_BATCH, client_name)

        # 校验是否
        validation_utils.check_dataframe_type(df)
        if (model_uri is None) == (local_uri is None):
            raise ValueError(
                "Either 'model_uri' or 'local_uri' must be provided, but not both."
            )
        if df.isStreaming:
            raise ValueError("Streaming DataFrames are not supported.")

        # 返回结果中会包含列名为 prediction，为预测结果，输入数据中不用此名字
        if PREDICTION_COLUMN_NAME in df.columns:
            raise ValueError(
                "FeatureStoreClient.score_batch returns a DataFrame with a new column "
                f'"{PREDICTION_COLUMN_NAME}". df already has a column with name '
                f'"{PREDICTION_COLUMN_NAME}".'
            )

        # 校验列中是否有重复列名
        validation_utils.validate_strings_unique(
            df.columns,
            "The provided DataFrame for scoring must have unique column names. Found duplicates {}.",
        )
        artifact_path = os.path.join("artifacts", MODEL_DATA_PATH_ROOT)
        print(f"artifact_path: {artifact_path}")
        with (TempDir() as tmp_location):
            local_path = (
                local_uri
                if local_uri
                else common_utils.download_model_artifacts(model_uri, tmp_location.path())
            )
            print(f"wedata local_path:{local_path}")
            model_data_path = os.path.join(local_path, artifact_path)
            print(f"artifact_path: {artifact_path}")

        # Augment local workspace metastore tables from 2L to 3L,
            # this will prevent us from erroneously reading data from other catalogs
            feature_spec = uc_utils.get_feature_spec_with_full_table_names(
                FeatureSpec.load(model_data_path)
            )
            raw_model_path = os.path.join(
                model_data_path, constants.RAW_MODEL_FOLDER
            )
            print(f"raw_model_path: {raw_model_path}")
        # 构建 udf 函数
            predict_udf =self._spark_client.get_predict_udf(
                raw_model_path,
                result_type=result_type,
                env_manager=env_manager,
                params=params,
                prebuilt_env_uri=kwargs.get(_PREBUILT_ENV_URI, None))
            # TODO (ML-17260) Consider reading the timestamp from the backend instead of feature store artifacts
            ml_model = Model.load(
                os.path.join(local_path, constants.ML_MODEL)
            )
            model_creation_timestamp_ms = (
                uc_utils.utc_timestamp_ms_from_iso_datetime_string(
                    ml_model.utc_time_created
                )
            )

        # Validate that columns needed for joining feature tables exist and are not duplicates.
        feature_input_keys = []
        for fci in feature_spec.feature_column_infos:
            feature_input_keys.extend([k for k in fci.lookup_key])
        on_demand_input_names = uc_utils.get_unique_list_order(
            [
                input_name
                for odci in feature_spec.on_demand_column_infos
                for input_name in odci.input_bindings.values()
            ]
        )
        intermediate_inputs = set(feature_input_keys + on_demand_input_names)

        print(f"wedata feature_input_keys:{feature_input_keys}")

        source_data_names = [
            sdci.name for sdci in feature_spec.source_data_column_infos
        ]
        print(f"wedata source_data_names:{source_data_names}")

        feature_output_names = [
            fci.output_name for fci in feature_spec.feature_column_infos
        ]
        print(f"wedata feature_output_names:{feature_output_names}")

        on_demand_output_names = [
            odci.output_name for odci in feature_spec.on_demand_column_infos
        ]
        all_output_names = set(
            source_data_names + feature_output_names + on_demand_output_names
        )
        print(f"wedata all_output_names:{all_output_names}")

        required_cols = intermediate_inputs.difference(all_output_names)
        required_cols.update(source_data_names)

        missing_required_columns = [
            col for col in required_cols if col not in df.columns
        ]
        if missing_required_columns:
            missing_columns_formatted = ", ".join(
                [f"'{s}'" for s in missing_required_columns]
            )
            raise ValueError(
                f"DataFrame is missing required columns {missing_columns_formatted}."
            )

        table_names = {fci.table_name for fci in feature_spec.feature_column_infos}
        feature_table_features_map = training_set_utils.get_features_for_tables(
            self._spark_client, table_names=table_names
        )
        feature_table_metadata_map = (
            training_set_utils.get_feature_table_metadata_for_tables(
                self._spark_client,
                table_names=table_names,
            )
        )
        feature_table_data_map = training_set_utils.load_feature_data_for_tables(
            self._spark_client, table_names=table_names
        )

        training_set_utils.validate_feature_column_infos_data(
            self._spark_client_helper,
            feature_spec.feature_column_infos,
            feature_table_features_map,
            feature_table_data_map,
        )

        # Check if the fetched feature tables match the feature tables logged in training
        self._warn_if_tables_mismatched_for_model(
            feature_spec=feature_spec,
            feature_table_metadata_map=feature_table_metadata_map,
            model_creation_timestamp_ms=model_creation_timestamp_ms,
        )

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

        # Required source data and feature lookup keys have been validated to exist in `df`.
        # No additional validation is required before resolving FeatureLookups and applying FeatureFunctions.
        training_set_utils.warn_if_non_photon_for_native_spark(
            kwargs.get(_USE_SPARK_NATIVE_JOIN, False), self._spark_client
        )

        augmented_df = TrainingSet(
            feature_spec=feature_spec,
            df=df,
            labels=[],
            feature_table_metadata_map=feature_table_metadata_map,
            feature_table_data_map=feature_table_data_map,
            uc_function_infos=uc_function_infos,
            use_spark_native_join=kwargs.get(_USE_SPARK_NATIVE_JOIN, False),
        )._augment_df()

        # Only included FeatureSpec columns should be part of UDF inputs for scoring.
        # Note: extra `df` columns not in FeatureSpec should be preserved.
        udf_input_columns = [
            ci.output_name for ci in feature_spec.column_infos if ci.include
        ]

        # Apply predictions.
        df_with_predictions = augmented_df.withColumn(
            PREDICTION_COLUMN_NAME, predict_udf(struct(*udf_input_columns))
        )

        # Reorder `df_with_predictions` to include:
        # 1. Preserved `df` columns, in `df` column order.
        # 2. Computed model input columns, in `FeatureSpec` column order.
        # 3. Prediction column.
        output_column_order = (
            df.columns
            + [col for col in udf_input_columns if col not in df.columns]
            + [PREDICTION_COLUMN_NAME]
        )
        return_df = df_with_predictions.select(output_column_order)
        return return_df

    def _warn_if_tables_mismatched_for_model(
        self,
        feature_spec: FeatureSpec,
        feature_table_metadata_map: Dict[str, FeatureTable],
        model_creation_timestamp_ms: float,
    ):
        """
        Helper method to warn if feature tables were deleted and recreated after a model was logged.
        For newer FeatureSpec versions >=3, we can compare the FeatureSpec and current table ids.
        Otherwise, we compare the model and table creation timestamps.
        """
        # 1. Compare feature table ids
        # Check for feature_spec logged with client versions that supports table_infos
        if len(feature_spec.table_infos) > 0:
            # When feature_spec.yaml is parsed, FeatureSpec.load will assert
            # that the listed table names in input_tables match table names in input_columns.
            # The following code assumes this as invariant and only checks for the table IDs.
            mismatched_tables = []
            for table_info in feature_spec.table_infos:
                feature_table = feature_table_metadata_map[table_info.table_name]
                if feature_table and table_info.table_id != feature_table.table_id:
                    mismatched_tables.append(table_info.table_name)
            if len(mismatched_tables) > 0:
                plural = len(mismatched_tables) > 1
                _logger.warning(
                    f"Feature table{'s' if plural else ''} {', '.join(mismatched_tables)} "
                    f"{'were' if plural else 'was'} deleted and recreated after "
                    f"the model was trained. Model performance may be affected if the features "
                    f"used in scoring have drifted from the features used in training."
                )

        # 2. Compare model creation timestamp with feature table creation timestamps
        feature_tables_created_after_model = []
        for name, metadata in feature_table_metadata_map.items():
            if model_creation_timestamp_ms < metadata.creation_timestamp:
                feature_tables_created_after_model.append(name)
        if len(feature_tables_created_after_model) > 0:
            plural = len(feature_tables_created_after_model) > 1
            message = (
                f"Feature table{'s' if plural else ''} {', '.join(feature_tables_created_after_model)} "
                f"{'were' if plural else 'was'} created after the model was logged. "
                f"Model performance may be affected if the features used in scoring have drifted "
                f"from the features used in training."
            )
            _logger.warning(message)
