from typing import List
from uuid import uuid4

from spetlr.etl import EtlBase, Orchestrator, dataset_group
from spetlr.etl.loaders import SimpleLoader
from spetlr.etl.loaders.simple_loader import Appendable
from spetlr.etl.transformers import UnionTransformer

from .log_transformer import LogTransformer


class LogOrchestrator(Orchestrator):
    """
    An extension of the Orchestrator class with built-in support for logging using
    LogTransformers.

    The LogOrchestrator class is a specialized subclass of the Orchestrator designed to
    facilitate ETL (Extract, Transform, Load) processes while providing seamless
    integration with LogTransformers. LogTransformers are specialized transformers that
    measure various metrics on dataframes during the ETL process, enabling the
    collection of valuable insights and measurements. Setting dataset output keys on
    log transformers is not required.

    Arguments:
        handles (List[Appendable]): A list of appendable handles that determine the
            sink location of the output generated by LogTransformers.
        suppress_composition_warning (bool, optional):
            Whether to suppress warnings about potential changes in the composition of
            the ETL process. Defaults to False.

    Methods:
        step(etl: EtlBase) -> LogOrchestrator:
            Adds the etl step to the orchestrator. It has the aliases 'extract_from',
            "transform_with', and 'load_into'. The method behaves the same as the step
            method of the Orchestrator class.
        step_log(log_transformers: LogTransformer) -> LogOrchestrator:
            Adds LogTransformers to the orchestrator. It has the alias 'log_with'. This
            step should be used when wanting to add a logging step.
        etl(inputs: dataset_group) ->dataset_group:
            Executes the ETL process. It has the alias 'execute'.
    """

    def __init__(
        self,
        handles: List[Appendable],
        suppress_composition_warning=False,
    ):
        self.handles = handles
        super().__init__(suppress_composition_warning)
        self.log_transformers_output_keys = []

    def step_log(
        self,
        log_transformer: LogTransformer,
    ) -> "LogOrchestrator":
        self.steps.append(log_transformer)

        # the output key of the LogTransformer is added to a list that is later used to
        # automatically load the data. Data is loaded with the handles provided when
        # instantiating the class.
        transformer_output_key = log_transformer.dataset_output_key

        # this actually ensures that even if no output key is set on the LogTransformer
        # the execute method will still operate correctly
        transformer_output_key = transformer_output_key or str(uuid4())
        self.log_transformers_output_keys.append(transformer_output_key)

        return self

    log_with = step_log

    def step(self, etl: EtlBase) -> "LogOrchestrator":
        return super().step(etl)

    extract_from = step
    transform_with = step
    load_into = step

    def etl(self, inputs: dataset_group = None) -> dataset_group:
        # First check if any log transformers have been added.
        # Note that if no log transformers have been added the LogOrchestrator will
        # behave just like a normal orchestrator.
        if len(self.log_transformers_output_keys) != 0:
            # if only a single log transformer has been added its dataset output key is
            # stored in a variable. This can then be written to the handles provided
            # when the class was instantiated.
            if len(self.log_transformers_output_keys) == 1:
                dataset_key_to_log = self.log_transformers_output_keys

            # if more than one log transformers has been added then these need to be
            # unioned before writing them using the handle(s) provided
            if len(self.log_transformers_output_keys) > 1:
                dataset_key_to_log = "UnionTransformer"
                self.steps.append(
                    UnionTransformer(
                        allowMissingColumns=True,
                        dataset_input_keys=self.log_transformers_output_keys,
                        dataset_output_key=dataset_key_to_log,
                        consume_inputs=True,
                    )
                )

            # now the output from the log transformer(s) can be loaded
            for handle in self.handles:
                self.steps.append(
                    SimpleLoader(
                        handle=handle,
                        mode="append",
                        dataset_input_keys=dataset_key_to_log,
                    )
                )

        # finally execute the orchestrator as usual
        return super().etl(inputs)

    execute = etl
