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.

        """

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

    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
