import logging

from datetime import datetime
from psi_elt.abstract_runner import AbstractRunner, NoneType
from psi_elt.custom_loader import CustomLoader
from typeguard import typechecked
from typing import Any, Dict


@typechecked
class LoaderRunner(AbstractRunner):
    def __init__(
        self,
        loaders: Dict[str, Any],
        aws_data_wrangler: Any = None,
        logger: logging.Logger = None
    ):
        '''
        Initializes a LoaderRunner.

        Accepts a loaders dict, mapping loader name to loader class.
        For example:
        {
            "slater": SlaterLoader,
            "tep": TepLoader
        }
        '''
        super().__init__(jobs=loaders, logger=logger)
        assert aws_data_wrangler, 'An AWS Data Wrangler utility is required to initialize a LoaderRunner'
        self.__aws_data_wrangler = aws_data_wrangler

    def parse_event(self, event: Dict[str, Any]) -> Dict[str, Any]:
        '''Parses a loader event. Child classes may override this function to
        implement custom event parsing rules.

        #### Required event keys
           1. name : the name of the loader 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. input_options.s3_uri                 : Exact S3 uri of the data to read
          4. input_options.s3_path                : S3 uri of the S3 path to look for data to read
          5. output_options.s3_uri                : Exact S3 uri where data will be written
          6. output_options.s3_path               : S3 uri of the S3 path to write data
          7. output_options.write_to_s3           : boolean for whether to write to S3
          8. output_options.update_athena         : boolean for whether to update Athena resources
        '''
        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': True,
            'load_kwargs': {},
            'to_parquet_kwargs': {}
        }

        # Parse input options
        if isinstance(event.get('input_options'), dict):
            input_options: dict = event['input_options']
            parsed_event['load_kwargs'] = {
                k: v for k, v in input_options.items()
                if v is not None and k in [
                    's3_uri',
                    's3_path'
                ]
            }

        # Parse output options
        if isinstance(event.get('output_options'), dict):
            output_options: dict = event['output_options']
            parsed_event['to_parquet_kwargs'] = {
                k: v for k, v in output_options.items()
                if v is not None and k in [
                    'update_athena'
                ]
            }
            if 'write_to_s3' in output_options:
                parsed_event['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,
            'write_to_s3': bool,
            'load_kwargs': {
                's3_uri': [str, NoneType],
                's3_path': [str, NoneType]
            },
            'to_parquet_kwargs': {
                'update_athena': [bool, NoneType]
            }
        }
        try:
            self._validate_dict(parsed_event, schema)
        except AssertionError as e:
            self.logger.error('Invalid event')
            raise e

    def _run(
        self,
        loader: CustomLoader,
        is_test: bool,
        capture_time: datetime,
        write_to_s3: bool,
        load_kwargs: dict,
        to_parquet_kwargs: dict
    ):
        data = loader.load(is_test, load_kwargs)
        if write_to_s3:
            self.logger.info('Beginning write to S3')
            self.__aws_data_wrangler.to_parquet(
                data=data,
                capture_time=capture_time,
                source=loader.name,
                **to_parquet_kwargs
            )
