import logging

from datetime import datetime
from psi_elt.abstract_runner import AbstractRunner, NoneType
from psi_elt.custom_extractor import CustomExtractor
from typeguard import typechecked
from typing import Any, Dict


@typechecked
class ExtractorRunner(AbstractRunner):
    def __init__(
        self,
        extractors: Dict[str, Any],
        s3_loader: Any = None,
        logger: 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
        if isinstance(event.get('input_options'), dict):
            input_options: dict = event['input_options']
            parsed_event['extract_kwargs'] = {
                k: v for k, v in input_options.items()
                if v is not None and k in [
                    'account_id'
                ]
            }

        # Parse output options
        if isinstance(event.get('output_options'), dict):
            output_options: dict = event['output_options']
            parsed_event['s3_load_kwargs'] = {
                k: v for k, v in output_options.items()
                if v is not None and k in [
                    's3_uri',
                    's3_path'
                ]
            }
            if 'write_to_s3' in output_options:
                parsed_event['extract_kwargs']['write_to_s3'] = output_options['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
