import logging

from datetime import datetime
from psi_elt.abstract_runner import AbstractRunner
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)
        assert s3_loader, 'A S3 loader is required to initialize an ExtractorRunner'
        self.__s3_loader = s3_loader

    def parse_event(self, event: Dict[str, Any]) -> Dict[str, Any]:
        '''Parses an extractor event like:
        {
            'name'          : 'extractor-name',   # required
            'capture_time'  : '2022-01-01T06:00',
            'is_test'       : False,
            'write_to_s3'   : False,
            'account_id'    : '012345678901',
            'bucket_name'   : 'my-s3-bucket',
            'output_s3_path': 'my/s3/path'
        }
        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. write_to_s3    : boolean for whether or not to write to S3
           4. account_id     : AWS account ID to extract data from
           5. bucket_name    : S3 bucket to write data to
           6. output_s3_path : S3 path to write data to

        Returns a parsed event like:
        {
            'name'               : 'extractor-name',
            'capture_time'       : datetime(2022, 1, 1, hour=6, tzinfo=pytz.utc),
            'job_init_kwargs'    : {},
            'is_test'            : False,
            'write_to_s3'        : False,
            'extract_kwargs'     : {
                'account_id'     : '012345678901'
            },
            's3_load_kwargs'     : {
                'bucket_name'    : 'my-s3-bucket',
                'output_s3_path' : 'my/s3/path'
            }
        }
        Child classes may override this function to implement custom event parsing rules.
        '''
        parsed_event = {
            'name': event.get('name'),
            'capture_time': self.get_capture_time(event),
            'job_init_kwargs': {},
            'is_test': bool(event.get('is_test')),
            'write_to_s3': event.get('write_to_s3', False),
            'extract_kwargs': {},
            's3_load_kwargs': {}
        }
        if event.get('account_id'):
            parsed_event['extract_kwargs']['account_id'] = event.get('account_id')
        if event.get('bucket_name'):
            parsed_event['s3_load_kwargs']['bucket_name'] = event.get('bucket_name')
        if event.get('output_s3_path'):
            parsed_event['s3_load_kwargs']['output_s3_path'] = event.get('output_s3_path')
        return parsed_event

    def validate_parsed_event(self, parsed_event: Dict[str, Any]):
        '''Validates that a parsed event has the required format'''
        required_keys = [
            'name', 'capture_time', 'job_init_kwargs', 'is_test', 'write_to_s3', 'extract_kwargs',
            's3_load_kwargs'
        ]
        missing_keys = [k for k in required_keys if k not in parsed_event]
        extra_keys = [k for k in parsed_event if k not in required_keys]
        assert not missing_keys, f'Parsed event is missing keys: {missing_keys}'
        assert not extra_keys, f'Parsed event has extra keys: {extra_keys}'
        assert isinstance(parsed_event['name'], str), f"Invalid name, expected a string but got: {parsed_event['name']}"
        assert isinstance(parsed_event['capture_time'], datetime), f"Invalid capture_time, expected a datetime but got: {parsed_event['capture_time']}"
        assert isinstance(parsed_event['job_init_kwargs'], dict), f"Invalid job_init_kwargs, expected a dict but got: {parsed_event['job_init_kwargs']}"
        assert isinstance(parsed_event['is_test'], bool), f"Invalid is_test, expected a bool but got: {parsed_event['is_test']}"
        assert isinstance(parsed_event['write_to_s3'], bool), f"Invalid write_to_s3, expected a bool but got: {parsed_event['write_to_s3']}"
        assert isinstance(parsed_event['extract_kwargs'], dict), f"Invalid extract_kwargs, expected a dict but got: {parsed_event['extract_kwargs']}"
        assert isinstance(parsed_event['s3_load_kwargs'], dict), f"Invalid s3_load_kwargs, expected a dict but got: {parsed_event['s3_load_kwargs']}"
        
        extract_kwargs = parsed_event['extract_kwargs']
        if 'account_id' in extract_kwargs:
            assert type(extract_kwargs['account_id']) in [type(None), str], f"Invalid extract_kwargs.account_id, expected a string but got: {extract_kwargs['account_id']}"

        s3_load_kwargs = parsed_event['s3_load_kwargs']
        if 'bucket_name' in s3_load_kwargs:
            assert isinstance(s3_load_kwargs['bucket_name'], str), f"Invalid s3_load_kwargs.bucket_name, expected a string but got: {s3_load_kwargs['bucket_name']}"
        if 'output_s3_path' in s3_load_kwargs:
            assert isinstance(s3_load_kwargs['output_s3_path'], str), f"Invalid s3_load_kwargs.output_s3_path, expected a string but got: {s3_load_kwargs['output_s3_path']}"

    def _run(
        self,
        extractor: CustomExtractor,
        is_test: bool,
        capture_time: datetime,
        write_to_s3: bool,
        extract_kwargs: dict,
        s3_load_kwargs: dict
    ) -> Any:
        data = extractor.extract(is_test, extract_kwargs)
        s3_uri = None
        if write_to_s3:
            self.logger.info('Beginning upload to S3')
            s3_uri = self.__s3_loader.load(
                extractor.name,
                data,
                capture_time,
                **s3_load_kwargs
            )
        return s3_uri
