import importlib
import inspect
import os
import posixpath
from typing import Any, Dict, List, Optional
from uuid import uuid4

import cloudpickle as cp
import pandas as pd

from snowflake.ml._internal import telemetry
from snowflake.ml._internal.utils import (
    identifier,
    pkg_version_utils,
    snowpark_dataframe_utils,
)
from snowflake.ml._internal.utils.query_result_checker import SqlResultValidator
from snowflake.ml._internal.utils.temp_file_utils import (
    cleanup_temp_files,
    get_temp_file_path,
)
from snowflake.ml.modeling._internal.estimator_utils import handle_inference_result
from snowflake.snowpark import DataFrame, Session, functions as F, types as T
from snowflake.snowpark._internal.utils import (
    TempObjectType,
    random_name_for_temp_object,
)

cp.register_pickle_by_value(inspect.getmodule(get_temp_file_path))
cp.register_pickle_by_value(inspect.getmodule(identifier.get_inferred_name))
cp.register_pickle_by_value(inspect.getmodule(handle_inference_result))

_PROJECT = "ModelDevelopment"


def _get_rand_id() -> str:
    """
    Generate random id to be used in sproc and stage names.

    Returns:
        Random id string usable in sproc, table, and stage names.
    """
    return str(uuid4()).replace("-", "_").upper()


class SnowparkTransformHandlers:
    def __init__(
        self,
        dataset: DataFrame,
        estimator: object,
        class_name: str,
        subproject: str,
        autogenerated: Optional[bool] = False,
    ) -> None:
        """
        Args:
            dataset: The dataset to run transform functions on.
            estimator: The estimator used to run transforms.
            class_name: class name to be used in telemetry.
            subproject: subproject to be used in telemetry.
            autogenerated: Whether the class was autogenerated from a template.
        """
        self.dataset = dataset
        self.estimator = estimator
        self._class_name = class_name
        self._subproject = subproject
        self._autogenerated = autogenerated

    def batch_inference(
        self,
        inference_method: str,
        input_cols: List[str],
        expected_output_cols: List[str],
        session: Session,
        dependencies: List[str],
        drop_input_cols: Optional[bool] = False,
        expected_output_cols_type: Optional[str] = "",
        *args: Any,
        **kwargs: Any,
    ) -> DataFrame:
        """Run batch inference on the given dataset.

        Args:
            session: An active Snowpark Session.
            dependencies: List of dependencies for the transformer.
            inference_method: the name of the method used by `estimator` to run inference.
            input_cols: List of feature columns for inference.
            expected_output_cols: column names (in order) of the output dataset.
            drop_input_cols: Boolean to determine whether to drop the input columns from the output dataset.
            expected_output_cols_type: Expected type of the output columns.
            args: additional positional arguments.
            kwargs: additional keyword args.

        Returns:
            A new dataset of the same type as the input dataset.
        """

        dependencies = self._get_validated_snowpark_dependencies(session, dependencies)
        dataset = self.dataset
        estimator = self.estimator
        # Register vectorized UDF for batch inference
        batch_inference_udf_name = random_name_for_temp_object(TempObjectType.FUNCTION)

        dataset = snowpark_dataframe_utils.cast_snowpark_dataframe_column_types(dataset)
        # Align the input_cols with snowpark dataframe's column name
        # This step also makes sure that the every col in input_cols exists in the current dataset
        snowpark_cols = dataset.select(input_cols).columns

        # Infer the datatype from input dataset's schema for batch inference
        # This is required before registering the UDTF
        fields = dataset.select(input_cols).schema.fields
        input_datatypes = []
        for field in fields:
            input_datatypes.append(field.datatype)

        statement_params = telemetry.get_function_usage_statement_params(
            project=_PROJECT,
            subproject=self._subproject,
            function_name=telemetry.get_statement_params_full_func_name(inspect.currentframe(), self._class_name),
            api_calls=[F.pandas_udf],
            custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
        )

        @F.pandas_udf(  # type: ignore[arg-type, misc]
            is_permanent=False,
            name=batch_inference_udf_name,
            packages=dependencies,  # type: ignore[arg-type]
            replace=True,
            session=session,
            statement_params=statement_params,
            input_types=[T.PandasDataFrameType(input_datatypes)],
        )
        def vec_batch_infer(input_df: pd.DataFrame) -> T.PandasSeries[dict]:  # type: ignore[type-arg]
            import numpy as np  # noqa: F401
            import pandas as pd

            input_df.columns = snowpark_cols

            if hasattr(estimator, "n_jobs"):
                # Vectorized UDF cannot handle joblib multiprocessing right now, deactivate the n_jobs
                estimator.n_jobs = 1
            inference_res = getattr(estimator, inference_method)(input_df, *args, **kwargs)

            transformed_numpy_array, _ = handle_inference_result(
                inference_res=inference_res,
                output_cols=expected_output_cols,
                inference_method=inference_method,
                within_udf=True,
            )

            if len(transformed_numpy_array.shape) > 1:
                if transformed_numpy_array.shape[1] != len(expected_output_cols):
                    series = pd.Series(transformed_numpy_array.tolist())
                    transformed_pandas_df = pd.DataFrame(series, columns=expected_output_cols)
                else:
                    transformed_pandas_df = pd.DataFrame(transformed_numpy_array.tolist(), columns=expected_output_cols)
            else:
                transformed_pandas_df = pd.DataFrame(transformed_numpy_array, columns=expected_output_cols)

            return transformed_pandas_df.to_dict("records")  # type: ignore[no-any-return]

        # Run Transform and get intermediate result
        INTERMEDIATE_OBJ_NAME = "tmp_result"
        # Use snowpark_cols can make sure the name ordering of the input dataframe
        # and only select those columns to put into vectorized udf
        output_obj = F.call_udf(batch_inference_udf_name, [F.col(col_name) for col_name in snowpark_cols])
        df_res: DataFrame = dataset.with_column(INTERMEDIATE_OBJ_NAME, output_obj)

        # Prepare the output
        output_cols = []
        output_col_names = []
        # When there is no expected_output_cols_type, default set it as StringType
        # snowpark cannot handle empty string, so this step give "string" value
        if expected_output_cols_type == "":
            expected_output_cols_type = "string"
        assert expected_output_cols_type is not None
        for output_feature in expected_output_cols:
            output_cols.append(F.col(INTERMEDIATE_OBJ_NAME)[output_feature].astype(expected_output_cols_type))
            output_col_names.append(identifier.get_inferred_name(output_feature))

        # Extract output from INTERMEDIATE_OBJ_NAME and drop that column
        df_res = df_res.with_columns(
            output_col_names,
            output_cols,
        ).drop(INTERMEDIATE_OBJ_NAME)

        if drop_input_cols:
            df_res = df_res.drop(*input_cols)

        return df_res

    def score(
        self,
        input_cols: List[str],
        label_cols: List[str],
        session: Session,
        dependencies: List[str],
        score_sproc_imports: List[str],
        sample_weight_col: Optional[str] = None,
        *args: Any,
        **kwargs: Any,
    ) -> float:
        """Score the given test dataset.

        Args:
            session: An active Snowpark Session.
            dependencies: score function dependencies.
            score_sproc_imports: imports for score stored procedure.
            input_cols: List of feature columns for inference.
            label_cols: List of label columns for scoring.
            sample_weight_col: A column assigning relative weights to each row for scoring.
            args: additional positional arguments.
            kwargs: additional keyword args.

        Returns:
            An accuracy score for the model on the given test data.
        """
        dependencies = self._get_validated_snowpark_dependencies(session, dependencies)
        dependencies.append("snowflake-snowpark-python")
        dataset = self.dataset
        estimator = self.estimator
        dataset = snowpark_dataframe_utils.cast_snowpark_dataframe_column_types(dataset)

        # Extract queries that generated the dataframe. We will need to pass it to score procedure.
        queries = dataset.queries["queries"]

        # Create a temp file and dump the score to that file.
        local_score_file_name = get_temp_file_path()
        with open(local_score_file_name, mode="w+b") as local_score_file:
            cp.dump(estimator, local_score_file)

        # Create temp stage to run score.
        score_stage_name = random_name_for_temp_object(TempObjectType.STAGE)
        assert session is not None  # keep mypy happy
        stage_creation_query = f"CREATE OR REPLACE TEMPORARY STAGE {score_stage_name};"
        SqlResultValidator(session=session, query=stage_creation_query).has_dimensions(
            expected_rows=1, expected_cols=1
        ).validate()

        # Use posixpath to construct stage paths
        stage_score_file_name = posixpath.join(score_stage_name, os.path.basename(local_score_file_name))
        score_sproc_name = random_name_for_temp_object(TempObjectType.PROCEDURE)
        statement_params = telemetry.get_function_usage_statement_params(
            project=_PROJECT,
            subproject=self._subproject,
            function_name=telemetry.get_statement_params_full_func_name(
                inspect.currentframe(), self.__class__.__name__
            ),
            api_calls=[F.sproc],
            custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
        )
        # Put locally serialized score on stage.
        session.file.put(
            local_score_file_name,
            stage_score_file_name,
            auto_compress=False,
            overwrite=True,
            statement_params=statement_params,
        )

        @F.sproc(  # type: ignore[misc]
            is_permanent=False,
            name=score_sproc_name,
            packages=dependencies,  # type: ignore[arg-type]
            replace=True,
            session=session,
            statement_params=statement_params,
            anonymous=True,
        )
        def score_wrapper_sproc(
            session: Session,
            sql_queries: List[str],
            stage_score_file_name: str,
            input_cols: List[str],
            label_cols: List[str],
            sample_weight_col: Optional[str],
            score_statement_params: Dict[str, str],
        ) -> float:
            import inspect
            import os

            import cloudpickle as cp

            for import_name in score_sproc_imports:
                importlib.import_module(import_name)

            for query in sql_queries[:-1]:
                _ = session.sql(query).collect(statement_params=score_statement_params)
            sp_df = session.sql(sql_queries[-1])
            df: pd.DataFrame = sp_df.to_pandas(statement_params=score_statement_params)
            df.columns = sp_df.columns

            local_score_file_name = get_temp_file_path()
            session.file.get(stage_score_file_name, local_score_file_name, statement_params=score_statement_params)

            local_score_file_name_path = os.path.join(local_score_file_name, os.listdir(local_score_file_name)[0])
            with open(local_score_file_name_path, mode="r+b") as local_score_file_obj:
                estimator = cp.load(local_score_file_obj)

            argspec = inspect.getfullargspec(estimator.score)
            if "X" in argspec.args:
                args = {"X": df[input_cols]}
            elif "X_test" in argspec.args:
                args = {"X_test": df[input_cols]}
            else:
                raise RuntimeError("Neither 'X' or 'X_test' exist in argument")

            if label_cols:
                label_arg_name = "Y" if "Y" in argspec.args else "y"
                args[label_arg_name] = df[label_cols].squeeze()

            if sample_weight_col is not None and "sample_weight" in argspec.args:
                args["sample_weight"] = df[sample_weight_col].squeeze()

            result: float = estimator.score(**args)
            return result

        # Call score sproc
        score_statement_params = telemetry.get_function_usage_statement_params(
            project=_PROJECT,
            subproject=self._subproject,
            function_name=telemetry.get_statement_params_full_func_name(
                inspect.currentframe(), self.__class__.__name__
            ),
            api_calls=[Session.call],
            custom_tags=dict([("autogen", True)]) if self._autogenerated else None,
        )

        kwargs = telemetry.get_sproc_statement_params_kwargs(score_wrapper_sproc, score_statement_params)
        score: float = score_wrapper_sproc(
            session,
            queries,
            stage_score_file_name,
            input_cols,
            label_cols,
            sample_weight_col,
            score_statement_params,
            **kwargs,
        )

        cleanup_temp_files([local_score_file_name])

        return score

    def _get_validated_snowpark_dependencies(self, session: Session, dependencies: List[str]) -> List[str]:
        """A helper function to validate dependencies and return the available packages that exists
        in the snowflake anaconda channel

        Args:
            session: the active snowpark Session
            dependencies: unvalidated dependencies

        Returns:
            A list of packages present in the snoflake conda channel.
        """

        return pkg_version_utils.get_valid_pkg_versions_supported_in_snowflake_conda_channel(
            pkg_versions=dependencies, session=session, subproject=self._subproject
        )
