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

from typeguard import typechecked

from psi_elt.abstract_runner import AbstractRunner, NoneType
from psi_elt.custom_extractor import CustomExtractor
from psi_elt.utils import filter_dict


@typechecked
class ExtractorRunner(AbstractRunner):
    def __init__(
        self,
        extractors: dict[str, Any],
        s3_loader: Optional[Any] = None,
        logger: Optional[logging.Logger] = None,
    ):
        """
        Initializes a ExtractorRunner.

        Accepts a extractors dict, mapping extractor name to extractor class.
        For example:
        {
            "divvy": DivvyExtractor,
            "pagerduty": PagerDutyExtractor
        }
        """
        super().__init__(jobs=extractors, logger=logger)
        self.__s3_loader = s3_loader

    def parse_event(self, event: dict[str, Any]) -> dict[str, Any]:
        """Parses an extractor event. Child classes may override this function to
        implement custom event parsing rules.

        #### Required event keys
           1. name : the name of the extractor to run

        #### Optional event keys
           1. capture_time               : target time for the job to run against
           2. is_test                    : boolean for whethor or not this is a test run
           3. input_options.account_id   : AWS account ID to extract data from
           4. output_options.s3_uri      : S3 location to write data to
           5. output_options.s3_path     : S3 path to write data to
           6. output_options.write_to_s3 : boolean for whether or not to write to S3
        """
        parsed_event = {
            "name": event.get("name"),
            "capture_time": self.get_capture_time(event),
            "job_init_kwargs": {},
            "is_test": bool(event.get("is_test")),
            "extract_kwargs": {},
            "s3_load_kwargs": {},
        }

        # Parse input options
        parsed_event["extract_kwargs"] = filter_dict(
            dct=event.get("input_options"), keys=["account_id"]
        )

        # Parse output options
        parsed_event["s3_load_kwargs"] = filter_dict(
            dct=event.get("output_options"), keys=["s3_uri", "s3_path"]
        )
        parsed_event["extract_kwargs"].update(
            filter_dict(dct=event.get("output_options"), keys=["write_to_s3"])
        )
        return parsed_event

    def validate_parsed_event(self, parsed_event: dict[str, Any]):
        """Validates that a parsed event has the required format"""
        schema = {
            "name": str,
            "capture_time": datetime,
            "job_init_kwargs": {},
            "is_test": bool,
            "extract_kwargs": {"account_id": [str, NoneType], "write_to_s3": [bool, NoneType]},
            "s3_load_kwargs": {"s3_uri": [str, NoneType], "s3_path": [str, NoneType]},
        }
        try:
            self._validate_dict(parsed_event, schema)
        except AssertionError as e:
            self.logger.error("Invalid event")
            raise e

    def _run(
        self,
        name: str,
        capture_time: datetime,
        job_init_kwargs: dict,
        is_test: bool,
        extract_kwargs: dict,
        s3_load_kwargs: dict,
    ) -> CustomExtractor:
        job_init_kwargs["s3_loader"] = self.__s3_loader
        extractor: CustomExtractor = self.get_job(name, capture_time, job_init_kwargs)
        extractor.extract(is_test, extract_kwargs, s3_load_kwargs)
        return extractor
