from typing import List, Optional

from snowflake.snowpark import DataFrame, Session


class MLRuntimeModelTrainer:
    """ML model training using the ml runties client."""

    def __init__(
        self,
        estimator: object,
        dataset: DataFrame,
        session: Session,
        input_cols: List[str],
        label_cols: Optional[List[str]],
        sample_weight_col: Optional[str],
        autogenerated: bool = False,
        subproject: str = "",
    ) -> None:
        """
        Initializes the MLRuntimeModelTrainer with a model, a Snowpark DataFrame, feature, and label column names.

        Args:
            estimator: SKLearn compatible estimator or transformer object.
            dataset: The dataset used for training the model.
            session: Snowflake session object to be used for training.
            input_cols: The name(s) of one or more columns in a DataFrame containing a feature to be used for training.
            label_cols: The name(s) of one or more columns in a DataFrame representing the target variable(s) to learn.
            sample_weight_col: The column name representing the weight of training examples.
            autogenerated: A boolean denoting if the trainer is being used by autogenerated code or not.
            subproject: subproject name to be used in telemetry.

        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.estimator = estimator
        self.dataset = dataset
        self.session = session
        self.input_cols = input_cols
        self.label_cols = label_cols
        self.sample_weight_col = sample_weight_col
        self._autogenerated = autogenerated
        self._subproject = subproject
        self._class_name = estimator.__class__.__name__

    def train(self) -> object:
        """
        Trains the model by pushing down the compute into SPCS ML Runtime
        """
        return self.client.train(
            estimator=self.estimator,
            dataset=self.dataset,
            input_cols=self.input_cols,
            label_cols=self.label_cols,
            sample_weight_col=self.sample_weight_col,
        )
