# Copyright (c) 2020 CloudZero - ALL RIGHTS RESERVED - PROPRIETARY AND CONFIDENTIAL
# Unauthorized copying of this file and/or project, via any medium is strictly prohibited.
# Direct all questions to legal@cloudzero.com
from typing import Any, List

import boto3
import pyfaaster.aws.handlers_decorators_v2 as faaster

from czc import json
from czc.constants import FEATURE_NAME

s3_client = boto3.client('s3')
sts_client = boto3.client('sts')

DEFAULT_URL_EXPIRATION = 900


@faaster.namespace_aware
def get_bucket_name(NAMESPACE):
    account_id = sts_client.get_caller_identity()['Account']
    return f'cz-{NAMESPACE}-{FEATURE_NAME}-{account_id}'


def get_presigned_url(bucket: str, key: str, expires_in: int = DEFAULT_URL_EXPIRATION) -> str:
    return s3_client.generate_presigned_url(
        'get_object',
        Params={
            'Bucket': bucket,
            'Key': key,
        },
        ExpiresIn=expires_in,
        HttpMethod='GET',
    )


def get_json_object(bucket: str, key: str) -> Any:
    response = s3_client.get_object(Bucket=bucket, Key=key)
    data = response['Body']
    return json.loads(data.read())


def put_json_object(bucket: str, key: str, data: Any) -> None:
    body = json.dumps(data).encode()
    s3_client.put_object(Bucket=bucket, Key=key, Body=body)


def list_objects(bucket: str, path: str) -> List[str]:
    page_iterator = _get_list_objects_page_iterator(bucket, Prefix=path)
    return [obj['Key'] for page in page_iterator for obj in page.get('Contents', [])]


def list_common_prefixes(bucket: str, prefix: str, delimiter: str) -> List[str]:
    page_iterator = _get_list_objects_page_iterator(bucket, Prefix=prefix, Delimiter=delimiter)
    return [obj['Prefix'] for page in page_iterator for obj in page.get('CommonPrefixes', [])]


def _get_list_objects_page_iterator(bucket: str, **kwargs):
    paginator = s3_client.get_paginator('list_objects_v2')
    return paginator.paginate(Bucket=bucket, **kwargs)
