import json
from datetime import datetime
from typing import Optional, Union

import boto3
from botocore.exceptions import ClientError

from psi_utils.logger import Logger
from psi_utils.json_utils import JsonEncoder


class S3Loader:
    def __init__(
        self,
        bucket: Optional[str] = None,
        path: Optional[str] = None,
        aws_session: Optional[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 Logger("psi-s3-loader").get_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
