import logging
from datetime import datetime
from typing import Any, Optional

from typeguard import typechecked

from psi_elt.abstract_elt_job import AbstractELTJob, LOADER_ELT_JOB_TYPE


@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: Optional[Any] = None,
        aws_session: Optional[Any] = None,
        logger: Optional[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")
