import importlib
import logging
import os
import sys
import time

import cloudpickle
from packaging.version import Version

from mlflow.recipes.artifacts import DataframeArtifact, TransformerArtifact
from mlflow.recipes.cards import BaseCard
from mlflow.recipes.step import BaseStep
from mlflow.recipes.step import StepClass
from mlflow.recipes.utils.execution import get_step_output_path
from mlflow.recipes.utils.step import get_pandas_data_profiles, validate_classification_config
from mlflow.recipes.utils.tracking import get_recipe_tracking_config, TrackingConfig
from mlflow.exceptions import MlflowException, INVALID_PARAMETER_VALUE

_logger = logging.getLogger(__name__)

_USER_DEFINED_TRANSFORM_STEP_MODULE = "steps.transform"


def _generate_feature_names(num_features):
    max_length = len(str(num_features))
    return ["f_" + str(i).zfill(max_length) for i in range(num_features)]


def _get_output_feature_names(transformer, num_features, input_features):
    import sklearn

    # `get_feature_names_out` was introduced in scikit-learn 1.0.0.
    if Version(sklearn.__version__) < Version("1.0.0"):
        return _generate_feature_names(num_features)

    try:
        # `get_feature_names_out` fails if `transformer` contains a transformer that doesn't
        # implement `get_feature_names_out`. For example, `FunctionTransformer` only implements
        # `get_feature_names_out` when it's instantiated with `feature_names_out`.
        # In scikit-learn >= 1.1.0, all transformers implement `get_feature_names_out`.
        # In scikit-learn == 1.0.*, some transformers implement `get_feature_names_out`.
        return transformer.get_feature_names_out(input_features)
    except Exception as e:
        _logger.warning(
            f"Failed to get output feature names with `get_feature_names_out`: {e}. "
            "Falling back to using auto-generated feature names."
        )
        return _generate_feature_names(num_features)


def _validate_user_code_output(transformer_fn):
    transformer = transformer_fn()
    if transformer is not None and not (hasattr(transformer, "fit") and callable(transformer.fit)):
        raise MlflowException(
            message="The transformer provided doesn't have a fit method."
        ) from None

    if transformer is not None and not (
        hasattr(transformer, "transform") and callable(transformer.transform)
    ):
        raise MlflowException(
            message="The transformer provided doesn't have a transform method."
        ) from None

    return transformer


class TransformStep(BaseStep):
    def __init__(self, step_config, recipe_root):  # pylint: disable=useless-super-delegation
        super().__init__(step_config, recipe_root)
        self.tracking_config = TrackingConfig.from_dict(self.step_config)

    def _validate_and_apply_step_config(self):
        self.target_col = self.step_config.get("target_col")
        self.positive_class = self.step_config.get("positive_class")
        if self.target_col is None:
            raise MlflowException(
                "Missing target_col config in recipe config.",
                error_code=INVALID_PARAMETER_VALUE,
            )
        if "using" in self.step_config:
            if self.step_config["using"] not in ["custom"]:
                raise MlflowException(
                    f"Invalid transform step configuration value {self.step_config['using']} for "
                    f"key 'using'. Supported values are: ['custom']",
                    error_code=INVALID_PARAMETER_VALUE,
                )
        else:
            self.step_config["using"] = "custom"
        self.run_end_time = None
        self.execution_duration = None
        self.skip_data_profiling = self.step_config.get("skip_data_profiling", False)

    def _run(self, output_directory):
        import pandas as pd

        run_start_time = time.time()

        train_data_path = get_step_output_path(
            recipe_root_path=self.recipe_root,
            step_name="split",
            relative_path="train.parquet",
        )
        train_df = pd.read_parquet(train_data_path)
        validate_classification_config(self.task, self.positive_class, train_df, self.target_col)

        validation_data_path = get_step_output_path(
            recipe_root_path=self.recipe_root,
            step_name="split",
            relative_path="validation.parquet",
        )
        validation_df = pd.read_parquet(validation_data_path)

        sys.path.append(self.recipe_root)

        def get_identity_transformer():
            from sklearn.pipeline import Pipeline
            from sklearn.preprocessing import FunctionTransformer

            return Pipeline(steps=[("identity", FunctionTransformer())])

        if "transformer_method" not in self.step_config and self.step_config["using"] == "custom":
            raise MlflowException(
                "Missing 'transformer_method' configuration in the transform step, "
                "which is using 'custom'.",
                error_code=INVALID_PARAMETER_VALUE,
            )
        method_config = self.step_config.get("transformer_method")
        transformer = None
        if method_config and self.step_config["using"] == "custom":
            transformer_fn = getattr(
                importlib.import_module(_USER_DEFINED_TRANSFORM_STEP_MODULE), method_config
            )
            transformer = _validate_user_code_output(transformer_fn)
        transformer = transformer if transformer else get_identity_transformer()
        transformer.fit(train_df.drop(columns=[self.target_col]), train_df[self.target_col])

        def transform_dataset(dataset):
            features = dataset.drop(columns=[self.target_col])
            transformed_features = transformer.transform(features)
            if not isinstance(transformed_features, pd.DataFrame):
                num_features = transformed_features.shape[1]
                columns = _get_output_feature_names(transformer, num_features, features.columns)
                transformed_features = pd.DataFrame(transformed_features, columns=columns)
            transformed_features[self.target_col] = dataset[self.target_col].values
            return transformed_features

        train_transformed = transform_dataset(train_df)
        validation_transformed = transform_dataset(validation_df)

        with open(os.path.join(output_directory, "transformer.pkl"), "wb") as f:
            cloudpickle.dump(transformer, f)

        train_transformed.to_parquet(
            os.path.join(output_directory, "transformed_training_data.parquet")
        )
        validation_transformed.to_parquet(
            os.path.join(output_directory, "transformed_validation_data.parquet")
        )

        self.run_end_time = time.time()
        self.execution_duration = self.run_end_time - run_start_time

        return self._build_profiles_and_card(train_df, train_transformed, transformer)

    def _build_profiles_and_card(self, train_df, train_transformed, transformer) -> BaseCard:
        # Build card
        card = BaseCard(self.recipe_name, self.name)

        if not self.skip_data_profiling:
            # Tab 1: build profiles for train_transformed
            train_transformed_profile = get_pandas_data_profiles(
                [["Profile of Train Transformed Dataset", train_transformed]]
            )
            card.add_tab("Data Profile (Train Transformed)", "{{PROFILE}}").add_pandas_profile(
                "PROFILE", train_transformed_profile
            )

        # Tab 3: transformer diagram
        from sklearn.utils import estimator_html_repr
        from sklearn import set_config

        set_config(display="diagram")
        transformer_repr = estimator_html_repr(transformer)
        card.add_tab("Transformer", "{{TRANSFORMER}}").add_html("TRANSFORMER", transformer_repr)

        # Tab 4: transformer input schema
        card.add_tab("Input Schema", "{{INPUT_SCHEMA}}").add_html(
            "INPUT_SCHEMA",
            BaseCard.render_table({"Name": n, "Type": t} for n, t in train_df.dtypes.items()),
        )

        # Tab 5: transformer output schema
        try:
            card.add_tab("Output Schema", "{{OUTPUT_SCHEMA}}").add_html(
                "OUTPUT_SCHEMA",
                BaseCard.render_table(
                    {"Name": n, "Type": t} for n, t in train_transformed.dtypes.items()
                ),
            )
        except Exception as e:
            card.add_tab("Output Schema", "{{OUTPUT_SCHEMA}}").add_html(
                "OUTPUT_SCHEMA", f"Failed to extract transformer schema. Error: {e}"
            )

        # Tab 6: transformer output data preview
        card.add_tab("Data Preview", "{{DATA_PREVIEW}}").add_html(
            "DATA_PREVIEW", BaseCard.render_table(train_transformed.head())
        )

        # Tab 7: run summary
        (
            card.add_tab(
                "Run Summary",
                """
                {{ EXE_DURATION }}
                {{ LAST_UPDATE_TIME }}
                """,
            )
        )

        return card

    @classmethod
    def from_recipe_config(cls, recipe_config, recipe_root):
        step_config = {}
        if recipe_config.get("steps", {}).get("transform", {}) is not None:
            step_config.update(recipe_config.get("steps", {}).get("transform", {}))
        step_config["target_col"] = recipe_config.get("target_col")
        step_config["recipe"] = recipe_config.get("recipe", "regression/v1")
        if "positive_class" in recipe_config:
            step_config["positive_class"] = recipe_config.get("positive_class")
        step_config.update(
            get_recipe_tracking_config(
                recipe_root_path=recipe_root,
                recipe_config=recipe_config,
            ).to_dict()
        )
        return cls(step_config, recipe_root)

    @property
    def name(self):
        return "transform"

    def get_artifacts(self):
        return [
            DataframeArtifact(
                "transformed_training_data",
                self.recipe_root,
                self.name,
                "transformed_training_data.parquet",
            ),
            DataframeArtifact(
                "transformed_validation_data",
                self.recipe_root,
                self.name,
                "transformed_validation_data.parquet",
            ),
            TransformerArtifact(
                "transformer", self.recipe_root, self.name, self.tracking_config.tracking_uri
            ),
        ]

    def step_class(self):
        return StepClass.TRAINING
