import boto3
import json

from botocore.exceptions import ClientError
from datetime import datetime
from psi_utils import DEFAULT_LOGGER
from psi_utils.json_utils import JsonEncoder
from typing import Union


class S3Loader:
    def __init__(
        self,
        bucket: str = None,
        path: str = None,
        aws_session: boto3.Session = None,
        s3_client=None,
        logger=None
    ):
        self.__bucket = bucket or 'psi-scorecard-metrics-us-east-1-490126044637'
        self.__path = path or 'scorecards-metrics'
        self.__aws_session = aws_session or boto3.Session()
        self.__s3_client = s3_client or self.__aws_session.client('s3')
        self.__logger = logger or DEFAULT_LOGGER

    def load(self, collector_name: str, data: Union[dict, list], capture_time: datetime) -> str:
        '''Load data into S3. Returns the S3 uri if successful.'''
        if data:
            assert collector_name, 'Cannot load S3 data without a collector name'
            assert capture_time, 'Cannot load S3 data without a capture time'
            assert isinstance(data, list) or isinstance(data, dict), f'Failed to load S3 data. Expected data to be a dict or list, got: {type(data)}'
            capture_date = capture_time.date()
            capture_time_ms = self.__dt_to_ms(capture_time)
            key = f'{self.__path}/{collector_name}/{capture_date.year}/{capture_date.month}/{capture_date.day}/{capture_time_ms}'
            return self.__put_object(
                key=key,
                body=json.dumps(data, cls=JsonEncoder),
                bucket=self.__bucket
            )
        return None

    def __put_object(self, key: str, body: str, bucket: str) -> str:
        """Put object into an S3 bucket. Returns the S3 uri if successful"""
        try:
            self.__s3_client.put_object(
                Body=body,
                Bucket=bucket,
                Key=key,
                ACL='bucket-owner-full-control'
            )
            return f's3://{bucket}/{key}'
        except ClientError as e:
            self.__logger.error(f'Failed to put S3 object s3://{bucket}/{key} with exception: {e}')
            raise e

    def __dt_to_ms(self, dt: datetime) -> int:
        if dt:
            return int(round(dt.timestamp() * 1000))
        return -1
