from typing import List

import pyspark.sql.functions as F
from pyspark.sql import DataFrame
from pyspark.sql.types import IntegerType, StringType, StructField, StructType

from spetlr.etl.log import LogTransformer
from spetlr.spark import Spark


class NullLogTransformer(LogTransformer):
    """
    A subclass of LogTransformer that logs the number of null values in a column. Note
    that this transformer relies on two separate .count() calls, meaning it can be an
    expensive operation if the input dataframe(s) have many rows.

    Arguments:
        log_name (str): The name for the log entry.
        column_name (str, optional): The specific column name for which the
            measurement is taken.
        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
    """

    def __init__(
        self,
        *,
        log_name: str,
        column_name: str,
        dataset_input_keys: List[str] = None,
        consume_inputs: bool = True,
    ):
        super().__init__(
            log_name=log_name,
            dataset_input_keys=dataset_input_keys,
            consume_inputs=consume_inputs,
        )
        self.column_name = column_name

    def log(self, df: DataFrame) -> DataFrame:
        n_rows = df.count()
        n_nulls = df.filter(F.col(self.column_name).isNull()).count()
        percentage_of_nulls = round(n_nulls / n_rows, 4)

        percentage_of_nulls_formatted = "{:.4f}".format(percentage_of_nulls)

        schema = StructType(
            [
                StructField("Count", IntegerType(), True),
                StructField("NumberOfNulls", IntegerType(), True),
                StructField("PercentageOfNulls", StringType(), True),
                StructField("ColumnName", StringType(), True),
            ]
        )

        data = [
            (n_rows, n_nulls, percentage_of_nulls_formatted, self.column_name),
        ]

        return Spark.get().createDataFrame(schema=schema, data=data)
