import uuid
import warnings
from abc import abstractmethod
from datetime import datetime
from functools import reduce
from typing import List

import pyspark.sql.functions as F
from pyspark.sql import DataFrame
from pyspark.sql.types import StringType, TimestampType

from spetlr.etl import Transformer
from spetlr.etl.types import dataset_group


class LogTransformer(Transformer):
    """
    Base class for specific ETL transformers that log measurements on a DataFrame.

    The `LogTransformer` class serves as a foundation for implementing various
    transformers that take measurements on a DataFrame. These measurements can
    include quantitative and qualitative values, such as counting rows,
    counting nulls, or any other custom metric that the user wishes to log
    during the ETL process. A log transformer produces a dataframe with a single line
    of data (measurement). It is recommended to use the log transformer with the
    LogOrchestrator class, so the load process is automatically dealt with.

    Usage:
        The `LogTransformer` class should be subclassed to create specific
        transformers that measure and log various metrics during the ETL
        process. Implement the `log` and `log_many` methods to define the
        specific measurement logic.

    Arguments:
        log_name (str): The name for the log entry.
        dataset_input_keys (List[str], optional): A list of input dataset keys
            that this transformer will process. Defaults to None.
        dataset_output_key (str, optional): The key for the output dataset
            generated by this transformer. If not provided, a random UUID will
            be used as the default value. It is not required for when using the
            LogOrchestrator.
        consume_inputs (bool, optional): Flag to control dataset consuming behavior.
            Defaults to True

    Methods:
        log(df: DataFrame) -> DataFrame:
            Abstract method that should be implemented in the derived classes.
            This method takes a DataFrame and logs the desired measurement(s).

        log_many(datasets: dataset_group) -> DataFrame:
            Abstract method that should be implemented in the derived classes.
            This method takes a dataset of dataframes and logs the desired
            measurement(s). Defaults to logging each frame by its `log` method.
    """

    def __init__(
        self,
        *,
        log_name: str,
        dataset_input_keys: List[str] = None,
        consume_inputs: bool = True,
    ):
        super().__init__(
            dataset_input_keys=dataset_input_keys,
            dataset_output_key=uuid.uuid4(),  # to ensure unique log output key
            consume_inputs=consume_inputs,
        )
        self.log_name = log_name
        self.timestamp_now = datetime.utcnow()
        self.log_method_name = self.__class__.__name__

    @abstractmethod
    def log(self, df: DataFrame) -> DataFrame:
        raise NotImplementedError()

    # This is the default method for log many, if the logging in the subclass can
    # be applied to each dataframe in the dataset separately. However, if the logging
    # logic should be performed across multiple dataframes in the dataset, then the
    # user should implement their own method.
    @abstractmethod
    def log_many(self, dataset: dataset_group) -> DataFrame:
        dfs_to_log = [
            self.log(df).withColumn("DatasetInputKey", F.lit(key))
            for key, df in dataset.items()
        ]

        df_log = reduce(
            lambda df1, df2: df1.unionByName(df2),
            dfs_to_log,
        )

        return df_log

    def _create_log(self, df_log: DataFrame) -> DataFrame:
        # some columns are casted to ensure the same types upon return

        df_log_cols = df_log.columns

        for col in [
            "LogName",
            "LogTimestamp",
            "LogMethodName",
            "LogId",
        ]:
            if col in df_log_cols:
                warnings.warn(
                    f"The Log Transformer used ({self.log_method_name}) produced a"
                    + f"column named '{col}'. Consider that this will be overwritten by"
                    + "the _create_log() method."
                )

        df = (
            df_log.withColumn("LogName", F.lit(self.log_name).cast(StringType()))
            .withColumn("LogTimestamp", F.lit(self.timestamp_now).cast(TimestampType()))
            .withColumn("LogMethodName", F.lit(self.log_method_name).cast(StringType()))
        )

        # this creates an id per log transformer
        df = df.withColumn("LogId", F.lit(str(uuid.uuid4())))

        columns_final = [
            "LogId",
            "LogName",
            "LogTimestamp",
            "LogMethodName",
        ] + df_log.columns

        df = df.select(*columns_final)

        return df

    def process(self, df: DataFrame) -> DataFrame:
        df = self.log(df)
        df = df.withColumn("DatasetInputKey", F.lit(self.dataset_input_keys[0]))
        df = self._create_log(df)
        return df

    def process_many(self, datasets: dataset_group) -> DataFrame:
        df = self.log_many(datasets)
        df = self._create_log(df)
        return df
