from typing import Any, List, Optional

from snowflake.ml.modeling._internal.snowpark_implementations.snowpark_handlers import (
    SnowparkTransformHandlers,
)
from snowflake.snowpark import DataFrame, Session


class MLRuntimeTransformHandlers:
    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.

        Raises:
            ModuleNotFoundError: The mlruntimes_client module is not available.
        """
        try:
            from snowflake.ml.runtime import MLRuntimeClient
        except ModuleNotFoundError as e:
            # This is an internal exception, not a user-facing one. The snowflake.ml.runtime module should
            # always be present when this class is instantiated.
            raise e

        self.client = MLRuntimeClient()
        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.
        Temporary workaround - pushdown implementation is not currently ready for batch_inference.
        We use a SnowparkTransformHandlers until we have a way to use the runtime client.

        Args:
            inference_method: the name of the method used by `estimator` to run inference.
            input_cols: List of feature columns for inference.
            session: An active Snowpark Session.
            dependencies: List of dependencies for the transformer.
            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.

        """

        mlrs_inference_methods = ["predict", "predict_proba", "predict_log_proba"]

        if inference_method in mlrs_inference_methods:
            result_df = self.client.inference(
                estimator=self.estimator,
                dataset=self.dataset,
                inference_method=inference_method,
                input_cols=input_cols,
                output_cols=expected_output_cols,
                drop_input_cols=drop_input_cols,
            )

        else:
            handler = SnowparkTransformHandlers(
                dataset=self.dataset,
                estimator=self.estimator,
                class_name=self._class_name,
                subproject=self._subproject,
                autogenerated=self._autogenerated,
            )
            result_df = handler.batch_inference(
                inference_method,
                input_cols,
                expected_output_cols,
                session,
                dependencies,
                drop_input_cols,
                expected_output_cols_type,
                *args,
                **kwargs,
            )

        assert isinstance(result_df, DataFrame)  # mypy - The MLRS return types are annotated as `object`.
        return result_df

    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.

        Raises:
            TypeError: The ML Runtimes client returned a non-float result
        """
        output_score = self.client.score(
            estimator=self.estimator,
            dataset=self.dataset,
            input_cols=input_cols,
            label_cols=label_cols,
            sample_weight_col=sample_weight_col,
        )
        if not isinstance(output_score, float):
            raise TypeError(
                f"The ML Runtimes Client returned a non-float value {output_score} of type {type(output_score)}"
            )
        return output_score
