import logging

from datetime import datetime
from psi_elt.abstract_elt_job import AbstractELTJob, LOADER_ELT_JOB_TYPE
from typeguard import typechecked
from typing import Any


@typechecked
class CustomLoader(AbstractELTJob):
    '''A CustomLoader is the second of the ELT jobs that is responsible for
    loading raw data into S3 in a queryable format.

    This class should be treated as abstract and should be extended.
    Child classes must override the _load function.
    '''

    def __init__(
        self,
        name: str,
        capture_time: datetime,
        aws_wrangler: Any = None,
        aws_session: Any = None,
        logger: logging.Logger = None
    ):
        super().__init__(
            name=name,
            elt_job_type=LOADER_ELT_JOB_TYPE,
            capture_time=capture_time,
            aws_session=aws_session,
            logger=logger
        )
        self.aws_wrangler = aws_wrangler

    def load(
        self,
        is_test: bool,
        load_kwargs: dict
    ):
        load_kwargs = self.remove_none_from_dict(load_kwargs)
        self._load(
            is_test,
            **load_kwargs
        )

    def _load(
        self,
        is_test: bool,
        **load_kwargs
    ) -> Any:
        raise Exception('CustomLoader._load must be overridden by the child class')
