import logging

from datetime import datetime
from psi_elt.abstract_runner import AbstractRunner
from psi_elt.custom_transformer import CustomTransformer
from typeguard import typechecked
from typing import Any, Dict


@typechecked
class TransformerRunner(AbstractRunner):
    def __init__(
        self,
        transformers: Dict[str, Any],
        timestream_loader: Any = None,
        s3_loader: Any = None,
        logger: logging.Logger = None
    ):
        '''
        Initializes a TransformerRunner.

        Accepts a transformers dict, mapping transformer name to transformer class.
        For example:
        {
            "tag-compliance": TagComplianceTransformer,
            "qualys": QualysTransformer
        }
        '''
        super().__init__(jobs=transformers, logger=logger)
        assert timestream_loader, 'A Timestream loader is required to initialize a TransformerRunner'
        assert s3_loader, 'A S3 loader is required to initialize a TransformerRunner'
        self.__timestream_loader = timestream_loader
        self.__s3_loader = s3_loader

    def parse_event(self, event: Dict[str, Any]) -> Dict[str, Any]:
        '''Parses a transformer event like:
        {
            'name'                  : 'transformer-name', # required
            'capture_time'          : '2022-01-01T06:00',
            'is_test'               : False,
            'write_to_s3'           : False,
            'write_to_timestream'   : False,
            's3_uri'                : 's3://bucket/path',
            'upsert'                : False,
            'max_batch_size'        : 50,
            'min_common_dimensions' : 2,
            'is_multi'              : False,           
            'bucket_name'           : 'my-s3-bucket',
            'output_s3_path'        : 'my/s3/path',
            's3_source'             : {
                's3_uri'            : 's3://bucket/path'
            }
        }
        Required event keys:
           1. name         : the name of the transformer to run

        Optional event keys:
           1. capture_time          : target time for the transformer 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. write_to_timestream   : boolean for whether or not to write to Timestream
           5. s3_uri                : S3 uri of an S3 object to retrieve data from
           6. upsert                : boolean for whether or not to upsert data (default True)
           7. max_batch_size        : int max batch size used when batching
           8. min_common_dimensions : int min common dimensions used when batching
           9. is_multi              : boolean for whether or not to use multi-measure records (default True)
          10. bucket_name           : S3 bucket to write data to
          11. output_s3_path        : S3 path to write data to
          12. s3_source             : Dictionary of S3 source info like an S3 uri. Overrides s3_uri.

        Returns a parsed event like:
        {
            'name'                      : 'transformer-name',
            'capture_time'              : datetime(2022, 1, 1, hour=6, tzinfo=pytz.utc),
            'job_init_kwargs'           : {},
            'is_test'                   : False,
            'write_to_s3'               : False,
            'write_to_timestream'       : False, 
            's3_uri'                    : 's3://bucket/path',
            'read_data_kwargs'          : {},
            'transform_kwargs'          : {},
            'timestream_load_kwargs'    : {
                'upsert':               : False,
                'ignore_parse_error'    : False,
                'max_batch_size'        : 50,
                'min_common_dimensions' : 2,
                'is_multi'              : False
            },
            '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),
            'write_to_timestream': event.get('write_to_timestream', True),
            's3_uri': event.get('s3_uri') or '',
            'read_data_kwargs': {},
            'transform_kwargs': {},
            'timestream_load_kwargs': {
                'upsert': bool(event.get('is_upsert', True)),
                'ignore_parse_error': False,
                'max_batch_size': event.get('max_batch_size'),
                'min_common_dimensions': event.get('min_common_dimensions'),
                'is_multi': bool(event.get('is_multi', True))
            },
            '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')
        if isinstance(event.get('s3_source'), dict):
            s3_source = event['s3_source']
            parsed_event['s3_uri'] = s3_source.get('s3_uri') or parsed_event['s3_uri']
        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',
            'write_to_timestream', 's3_uri', 'read_data_kwargs', 'transform_kwargs',
            'timestream_load_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['write_to_timestream'], bool), f"Invalid write_to_s3, expected a bool but got: {parsed_event['write_to_timestream']}"
        assert isinstance(parsed_event['s3_uri'], str), f"Invalid s3_uri, expected a string but got: {parsed_event['s3_uri']}"
        assert isinstance(parsed_event['read_data_kwargs'], dict), f"Invalid read_data_kwargs, expected a dict but got: {parsed_event['read_data_kwargs']}"
        assert isinstance(parsed_event['transform_kwargs'], dict), f"Invalid transform_kwargs, expected a dict but got: {parsed_event['transform_kwargs']}"
        assert isinstance(parsed_event['timestream_load_kwargs'], dict), f"Invalid timestream_load_kwargs, expected a dict but got: {parsed_event['timestream_load_kwargs']}"
        assert isinstance(parsed_event['s3_load_kwargs'], dict), f"Invalid s3_load_kwargs, expected a dict but got: {parsed_event['s3_load_kwargs']}"
        
        timestream_load_kwargs = parsed_event['timestream_load_kwargs']
        assert isinstance(timestream_load_kwargs.get('upsert'), bool), f"Invalid timestream_load_kwargs.upsert, expected a bool but got: {timestream_load_kwargs.get('upsert')}"
        assert isinstance(timestream_load_kwargs.get('ignore_parse_error'), bool), f"Invalid timestream_load_kwargs.ignore_parse_error, expected a bool but got: {timestream_load_kwargs.get('ignore_parse_error')}"
        assert 'max_batch_size' in timestream_load_kwargs, 'Parsed event missing key timestream_load_kwargs.max_batch_size'
        assert type(timestream_load_kwargs['max_batch_size']) in [int, type(None)], f"Invalid timestream_load_kwargs.max_batch_size, expected an int or None but got: {timestream_load_kwargs['max_batch_size']}"
        assert 'min_common_dimensions' in timestream_load_kwargs, 'Parsed event missing key timestream_load_kwargs.min_common_dimensions'
        assert type(timestream_load_kwargs['min_common_dimensions']) in [int, type(None)], f"Invalid timestream_load_kwargs.min_common_dimensions, expected an int or None but got: {timestream_load_kwargs['min_common_dimensions']}"
        assert isinstance(timestream_load_kwargs.get('is_multi'), bool), f"Invalid timestream_load_kwargs.is_multi, expected a bool but got: {timestream_load_kwargs.get('is_multi')}"

        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,
        transformer: CustomTransformer,
        is_test: bool,
        capture_time: datetime,
        s3_uri: str,
        write_to_s3: bool,
        write_to_timestream: bool,
        read_data_kwargs: dict,
        transform_kwargs: dict,
        timestream_load_kwargs: dict,
        s3_load_kwargs: dict
    ) -> Any:
        data = transformer.transform(is_test, s3_uri, read_data_kwargs, transform_kwargs)
        s3_uri = None
        if write_to_s3:
            self.logger.warning('Skipping upload to S3')
            # self.logger.info('Beginning upload to S3')
            # s3_uri = self.__s3_loader.load(
            #     transformer.name,
            #     data,
            #     capture_time,
            #     **s3_load_kwargs
            # )
        if write_to_timestream:
            self.logger.info('Beginning write to Timestream')
            result = self.__timestream_loader.load(
                data=data,
                capture_time=capture_time,
                source=transformer.name,
                **timestream_load_kwargs
            )
            if len(result['invalid_items']) > 0:
                raise Exception(f"Failed to write to Timestream with invalid items {result['invalid_items']}")
            if len(result['rejected_records']) > 0:
                raise Exception(f"Failed to write to Timestream with rejected records {result['rejected_records']}")
        return s3_uri
