import boto3
import logging
import pytz
import requests
import sys

from datetime import datetime
from psi_collector.custom_collector import CustomCollector
from typing import Any, Dict

class CollectorRunner():
    def __init__(
        self,
        collectors: Dict[str, Any],
        s3_loader: Any = None,
        timestream_loader: Any = None,
        logger: logging.Logger = None
    ):
        '''
        Initializes a CollectorRunner.

        Accepts a collectors dict, mapping collector name to collector class.
        For example:
        {
            "divvy": DivvyCollector,
            "pagerduty": PagerDutyCollector
        }
        
        '''
        assert collectors and isinstance(collectors, dict), 'A collectors dict is required to initialize a CollectorRunner'
        assert s3_loader, 'A S3 loader is required to initialize a CollectorRunner'
        assert timestream_loader, 'A Timestream loader is required to initialize a CollectorRunner'
        assert type(logger) in [logging.Logger, type(None)], f'A logger of type logging.Logger is required to initialize a CollectorRunner. Got {logger}'
        self.__collectors = collectors
        self.__s3_loader = s3_loader
        self.__timestream_loader = timestream_loader

        self.__logger = logger
        if not self.__logger:
            handler = logging.StreamHandler(sys.stdout)
            log = logging.getLogger('collector_runner')
            log.setLevel(logging.INFO)
            handler.setLevel(logging.INFO)
            formatter = logging.Formatter(
                '%(asctime)s [%(levelname)s] %(filename)s:%(lineno)d - %(message)s'
            )
            handler.setFormatter(formatter)
            log.addHandler(handler)
            log.propagate = False
            self.__logger = log

    def collect(self, event: dict):
        '''Initializes and runs a collector.

        Accepts an event dictionary with the following required attributes:
            1. collector_name : the name of the collector to run
            2. s3_path        : the full s3 path to the data to collect from S3

        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
        
        If writing to Timestream, the following attributes are required:
            1. write_to_timestream     : boolean that must be set to true if writing to Timestream
            2. database_name           : Timestream database to write to
            3. output_timestream_table : Timestream table to write to  
            (optional) is_upsert       : whether or not to upsert Timestream records (default True)
        
        Optionally, capture_time may be set on the event as a millisecond timestamp. If unset,
        capture_time is taken from either the s3_path or the current time.
        '''
        t0 = datetime.now()
        self.__collect(
            collector_name=event.get('collector_name'),
            capture_time=self.get_capture_time(event),
            write_to_s3=event.get('write_to_s3', True),
            input_s3_path=event.get('s3_path'),
            write_to_timestream=event.get('write_to_timestream', True),
            upsert=event.get('is_upsert', True),
            max_batch_size=event.get('max_batch_size'),
            min_common_dimensions=event.get('min_common_dimensions'),
            is_multi=event.get('is_multi'),
            account_id=event.get('account_id')
        )
        self.__logger.info(f'Collection completed in {str(datetime.now() - t0)}')

    def __collect(
        self,
        collector_name: str,
        capture_time: datetime,
        write_to_s3: bool,
        input_s3_path: str,
        write_to_timestream: bool,
        upsert: bool,
        max_batch_size: int,
        min_common_dimensions: int,
        is_multi: bool,
        account_id: str
    ):
        collector = self.get_collector(collector_name, capture_time=capture_time)
        data = collector.collect(s3_path=input_s3_path, account_id=account_id)

        if write_to_s3:
            self.__logger.info('Beginning upload to S3')
            self.__s3_loader.load(collector_name, data, capture_time)

        if write_to_timestream:
            self.__logger.info('Beginning write to Timestream')
            result = self.__timestream_loader.load(
                data=data,
                capture_time=capture_time,
                source=collector_name,
                upsert=upsert,
                ignore_parse_error=True,
                max_batch_size=max_batch_size,
                min_common_dimensions=min_common_dimensions,
                is_multi=is_multi
            )
            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']}")

    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_path' in event and event['s3_path']:
            capture_time_ms = event['s3_path'].split('/')[-1]
            # fromtimestamp returns a naive datetime object that IS adjusted for the local timezone
            return datetime.fromtimestamp(int(capture_time_ms)/1000).astimezone().astimezone(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_collector(
        self,
        collector_name: str,
        api=None,
        aws_session: boto3.Session = None,
        http_session: requests.Session = None,
        capture_time: datetime = None
    ) -> CustomCollector:
        assert capture_time, 'A capture_time is required to initialize a collector'
        cls = self.__collectors.get(collector_name)
        if cls:
            return cls(
                capture_time=capture_time,
                api=api,
                aws_session=aws_session,
                http_session=http_session
            )
        raise Exception(f'Unrecognized collector with name: {collector_name}')
