import logging
import pytz

from datetime import datetime
from psi_elt import get_logger
from psi_elt.abstract_elt_job import AbstractELTJob
from typeguard import typechecked
from typing import Any, Dict


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

        This class should be treated as abstract and should be extended.
        Child classes must override the _run function. Optionally, child
        classes may override the parse_event and validate_parsed_event
        functions as well.

        This class includes utilties for getting a capture time and
        initializing an ELT job.
        '''
        assert jobs, 'A jobs dict is required to initialize an AbstractRunner'
        self.__jobs = jobs
        self.__logger = logger or get_logger()

    @property
    def logger(self) -> logging.Logger:
        return self.__logger

    def run(self, event: Dict[str, Any]) -> Dict[str, Any]:
        t0 = datetime.now()
        parsed_event = self.parse_event(event)
        self.validate_parsed_event(parsed_event)
        job = self.get_job(
            job_name=parsed_event.pop('name'),
            capture_time=parsed_event['capture_time'],
            job_init_kwargs=parsed_event.pop('job_init_kwargs')
        )
        is_test = parsed_event.pop('is_test')
        s3_uri = self._run(job, is_test, **parsed_event)
        self.logger.info(f'ELT Job completed in {str(datetime.now() - t0)}')
        return {
            'event': event,
            'parsed_event': parsed_event,
            'job': job.to_dict(),
            'results': {
                's3_uri': s3_uri
            }
        }

    def _run(self, job: AbstractELTJob, **parsed_event):
        raise Exception('AbstractRunner._run should be overridden by the child class')

    def parse_event(self, event: Dict[str, Any]) -> Dict[str, Any]:
        '''Parses an ELT job event like:
        {
            'name'         : 'extractor-name',
            'capture_time' : '2022-01-01T06:00', # optional
            'is_test'      : False,              # optional
        }
        Required event keys:
           1. name : the name of the ELT job 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

        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
        }
        Child classes may override this function to implement custom event parsing rules.
        '''
        return {
            'name': event.get('name'),
            'capture_time': self.get_capture_time(event),
            'job_init_kwargs': {},
            'is_test': bool(event.get('is_test'))
        }

    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']
        missing_keys = [k for k in required_keys if k not in parsed_event]
        assert not missing_keys, f'Parsed event is missing keys: {missing_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']}"

    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)
        elif 's3_uri' in event and event['s3_uri']:
            capture_time_ms = event['s3_uri'].split('/')[-1]
            # fromtimestamp returns a naive datetime object that IS adjusted for the local timezone
            return datetime.fromtimestamp(int(capture_time_ms)/1000).astimezone(pytz.utc)
        else:
            # now returns a naive datetime object that IS adjusted for the local timezone
            return datetime.now().astimezone(pytz.utc)

    def get_job(
        self,
        job_name: str,
        capture_time: datetime,
        job_init_kwargs: dict
    ) -> AbstractELTJob:
        assert capture_time, 'A capture_time is required to initialize an ELT job'
        cls = self.__jobs.get(job_name)
        if cls:
            return cls(
                capture_time,
                **job_init_kwargs
            )
        raise Exception(f'Unrecognized ELT job with name: {job_name}')
