import boto3
import logging
import pytz
import re
import requests

from datetime import datetime
from psi_collector import DEFAULT_LOGGER, log_it
from psi_collector.custom_collector import CustomCollector
from typing import Any, Dict, Tuple

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 or DEFAULT_LOGGER

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

        Parameters
        ----------
        event : dict
            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.

        Returns
        -------
        result : dict
            Returns a collection result containing the data output location, e.g.:
                {
                    's3_uri': ...,
                    'timestream_db': ...,
                    'timestream_table': ...
                }
            Please note that either the s3_uri OR the timestream_db/table will be set,
            but not both.
        '''
        t0 = datetime.now()
        collector = self.get_collector(
            collector_name=event.get('collector_name'),
            capture_time=self.get_capture_time(event)
        )
        if event.get('s3_path'):
            timestream_db, timestream_table = self.__collect_from_s3(
                collector=collector,
                input_s3_path=event.get('s3_path'),
                write_to_timestream=event.get('write_to_timestream', True),
                upsert=event.get('is_upsert', True)
            )
            result = self.__format_result(
                timestream_db=timestream_db,
                timestream_table=timestream_table
            )
        else:
            s3_uri = self.__collect(
                collector=collector,
                write_to_s3=event.get('write_to_s3', True)
            )
            result = self.__format_result(s3_uri=s3_uri)
        self.__logger.info(f'Collection completed in {str(datetime.now() - t0)}')
        return result

    def __collect(self, collector: CustomCollector, write_to_s3: bool) -> str:
        '''This function calls collect on the collector and ONLY writes to S3'''
        data = collector.collect()
        if write_to_s3:
            self.__logger.info('Beginning upload to S3')
            return self.__s3_loader.load(collector.name, data, collector.capture_time)
        return None

    def __collect_from_s3(
        self,
        collector: CustomCollector,
        input_s3_path: str,
        write_to_timestream: bool,
        upsert: bool
    ) -> Tuple[str, str]:
        '''This function calls collect_from_s3 on the collector and ONLY writes to Timestream'''
        data = collector.collect_from_s3(input_s3_path)
        if write_to_timestream:
            self.__logger.info('Beginning write to Timestream')
            result = self.__timestream_loader.load(
                data=data,
                capture_time=collector.capture_time,
                source=collector.name,
                upsert=upsert,
                ignore_parse_error=True
            )
            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 self.__timestream_loader.database_name, self.__timestream_loader.table_name
        return None, None

    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}')

    def __format_result(
        self,
        s3_uri: str = None,
        timestream_db: str = None,
        timestream_table: str = None
    ):
        result = {'s3_uri': None, 'timestream_db': None, 'timestream_table': None}
        if s3_uri:
            assert not (timestream_db or timestream_table), 'A collection result cannot have both a S3 ouput and a Timestream output'
            assert re.match('^s3://([^/]+)/(.*?([^/]+)/?)$', s3_uri), f'Invalid s3 uri: {s3_uri}'
            result['s3_uri'] = str(s3_uri)
        elif timestream_db or timestream_table:
            assert timestream_db and timestream_table, 'A collection result Timestream output must include both a Timestream database and table'
            result['timestream_db'] = timestream_db
            result['timestream_table'] = timestream_table
        return result
