import boto3
import logging
import pytz
import requests

from datetime import datetime
from psi_extractor import get_logger
from psi_extractor.custom_extractor import CustomExtractor
from typeguard import typechecked
from typing import Any, Dict


@typechecked
class ExtractorRunner():
    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
        }
        
        '''
        assert extractors, 'An extractors dict is required to initialize a ExtractorRunner'
        assert s3_loader, 'A S3 loader is required to initialize a ExtractorRunner'
        self.__extractors = extractors
        self.__s3_loader = s3_loader
        self.__logger = logger or get_logger()

    def extract(self, event: Dict[str, Any]) -> datetime:
        '''Initializes and runs a extractor.

        Accepts an event dictionary with the following required attributes:
            1. extractor_name : the name of the extractor to run

        If writing to S3, the following attributes are required:
            1. write_to_s3    : boolean that must be set to true if writing to S3
            2. bucket_name    : S3 bucket to write to
            3. output_s3_path : S3 path to write to
        
        Optionally, capture_time may be set on the event as a millisecond timestamp. If unset,
        capture_time is set to the current time.
        '''
        t0 = datetime.now()
        parsed_event = self.parse_event(event)
        self.__validate_parsed_event(parsed_event)
        self.__extract(**parsed_event)
        self.__logger.info(f'Extraction completed in {str(datetime.now() - t0)}')
        return parsed_event['capture_time']

    def parse_event(self, event: Dict[str, Any]) -> Dict[str, Any]:
        parsed_event = {
            'extractor_name': event.get('extractor_name'),
            'capture_time': self.get_capture_time(event),
            'write_to_s3': event.get('write_to_s3', True),
            'extractor_init_kwargs': {},
            'extract_kwargs': {
                'account_id': event.get('account_id')
            },
            's3_load_kwargs': {}
        }
        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]):
        required_keys = [
            'extractor_name','capture_time','write_to_s3',
            'extractor_init_kwargs','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['extractor_name'], str), f"Invalid extractor_name, expected a string but got: {parsed_event['extractor_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['write_to_s3'], bool), f"Invalid write_to_s3, expected a bool but got: {parsed_event['write_to_s3']}"
        assert isinstance(parsed_event['extractor_init_kwargs'], dict), f"Invalid extractor_init_kwargs, expected a dict but got: {parsed_event['extractor_init_kwargs']}"
        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 __extract(
        self,
        extractor_name: str,
        capture_time: datetime,
        write_to_s3: bool,
        extractor_init_kwargs: dict,
        extract_kwargs: dict,
        s3_load_kwargs: dict
    ):
        extractor = self.get_extractor(extractor_name, capture_time, extractor_init_kwargs)
        data = extractor.extract(**extract_kwargs)
        if write_to_s3:
            self.__logger.info('Beginning upload to S3')
            self.__s3_loader.load(
                extractor_name,
                data,
                capture_time,
                **s3_load_kwargs
            )

    def get_capture_time(self, event: dict) -> datetime:
        if 'capture_time' in event and event['capture_time']:
            # fromisoformat returns a naive datetime object that IS NOT adjusted for the local timezone
            return datetime.fromisoformat(event['capture_time']).replace(tzinfo=pytz.utc)
        else:
            # now returns a naive datetime object that IS adjusted for the local timezone
            return datetime.now().astimezone().astimezone(pytz.utc)

    def get_extractor(
        self,
        extractor_name: str,
        capture_time: datetime,
        extractor_init_kwargs: dict
    ) -> CustomExtractor:
        assert capture_time, 'A capture_time is required to initialize an extractor'
        cls = self.__extractors.get(extractor_name)
        if cls:
            return cls(
                capture_time,
                **extractor_init_kwargs
            )
        raise Exception(f'Unrecognized extractor with name: {extractor_name}')
