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.mlflow_model import ModelWrapper
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,
    ):
        # 验证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

            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=ModelWrapper(model),
                #data_path=data_path,
                artifacts={"feature_store": data_path},
                code_path=None,
                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}")
