# implicit, not referenced - must be the 1st import
from ._logging_config import configure_logging
import logging
from ._config import Config
from ._timeslice import Timeslice
from ._stage_type import StageType


def yetl_flow(
    stage: StageType,
    project: str,
    pipeline: str = None,
    config_path: str = None,
):
    def decorate(function):
        def wrap_function(*args, **kwargs):
            configure_logging(project)
            _logger = logging.getLogger(__name__)

            _pipeline = pipeline
            if not _pipeline:
                _pipeline = function.__name__

            _logger.info(f"Loading pipeline configuration {_pipeline}")

            timeslice = kwargs.get("timeslice", Timeslice(day="*", month="*", year="*"))
            if "timeslice" in kwargs.keys():
                del kwargs["timeslice"]

            try:
                table = kwargs["table"]
                del kwargs["table"]
            except KeyError as e:
                raise Exception(f"{e} is a required argument for a yetl flow function")

            config = Config(
                project=project, pipeline=_pipeline, config_path=config_path
            )
            table_mapping = config.get_table_mapping(
                timeslice=timeslice, stage=stage, table=table
            )
            _logger.info(f"Calling function {function.__name__}")
            ret = function(
                *args,
                table_mapping=table_mapping,
                **kwargs,
            )
            return ret

        return wrap_function

    return decorate
