# Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
# https://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.
"""
This file contains private functionality for interacting with the AWS
Common Runtime library (awscrt) in boto3.

All code contained within this file is for internal usage within this
project and is not intended for external consumption. All interfaces
contained within are subject to abrupt breaking changes.
"""

import threading

import botocore.exceptions
from botocore.session import Session
from s3transfer.crt import (
    BotocoreCRTCredentialsWrapper,
    BotocoreCRTRequestSerializer,
    CRTTransferManager,
    acquire_crt_s3_process_lock,
    create_s3_crt_client,
)

# Singletons for CRT-backed transfers
CRT_S3_CLIENT = None
BOTOCORE_CRT_SERIALIZER = None

CLIENT_CREATION_LOCK = threading.Lock()
PROCESS_LOCK_NAME = 'boto3'


def _create_crt_client(session, config, region_name, cred_provider):
    """Create a CRT S3 Client for file transfer.

    Instantiating many of these may lead to degraded performance or
    system resource exhaustion.
    """
    create_crt_client_kwargs = {
        'region': region_name,
        'use_ssl': True,
        'crt_credentials_provider': cred_provider,
    }
    return create_s3_crt_client(**create_crt_client_kwargs)


def _create_crt_request_serializer(session, region_name):
    return BotocoreCRTRequestSerializer(
        session, {'region_name': region_name, 'endpoint_url': None}
    )


def _create_crt_s3_client(
    session, config, region_name, credentials, lock, **kwargs
):
    """Create boto3 wrapper class to manage crt lock reference and S3 client."""
    cred_wrapper = BotocoreCRTCredentialsWrapper(credentials)
    cred_provider = cred_wrapper.to_crt_credentials_provider()
    return CRTS3Client(
        _create_crt_client(session, config, region_name, cred_provider),
        lock,
        region_name,
        cred_wrapper,
    )


def _initialize_crt_transfer_primatives(client, config):
    lock = acquire_crt_s3_process_lock(PROCESS_LOCK_NAME)
    if lock is None:
        # If we're unable to acquire the lock, we cannot
        # use the CRT in this process and should default to
        # the classic s3transfer manager.
        return None, None

    session = Session()
    region_name = client.meta.region_name
    credentials = client._get_credentials()

    serializer = _create_crt_request_serializer(session, region_name)
    s3_client = _create_crt_s3_client(
        session, config, region_name, credentials, lock
    )
    return serializer, s3_client


def get_crt_s3_client(client, config):
    global CRT_S3_CLIENT
    global BOTOCORE_CRT_SERIALIZER

    with CLIENT_CREATION_LOCK:
        if CRT_S3_CLIENT is None:
            serializer, s3_client = _initialize_crt_transfer_primatives(
                client, config
            )
            BOTOCORE_CRT_SERIALIZER = serializer
            CRT_S3_CLIENT = s3_client

    return CRT_S3_CLIENT


class CRTS3Client:
    """
    This wrapper keeps track of our underlying CRT client, the lock used to
    acquire it and the region we've used to instantiate the client.

    Due to limitations in the existing CRT interfaces, we can only make calls
    in a single region and does not support redirects. We track the region to
    ensure we don't use the CRT client when a successful request cannot be made.
    """

    def __init__(self, crt_client, process_lock, region, cred_provider):
        self.crt_client = crt_client
        self.process_lock = process_lock
        self.region = region
        self.cred_provider = cred_provider


def is_crt_compatible_request(client, crt_s3_client):
    """
    Boto3 client must use same signing region and credentials
    as the CRT_S3_CLIENT singleton. Otherwise fallback to classic.
    """
    if crt_s3_client is None:
        return False

    boto3_creds = client._get_credentials()
    if boto3_creds is None:
        return False

    is_same_identity = compare_identity(
        boto3_creds.get_frozen_credentials(), crt_s3_client.cred_provider
    )
    is_same_region = client.meta.region_name == crt_s3_client.region
    return is_same_region and is_same_identity


def compare_identity(boto3_creds, crt_s3_creds):
    try:
        crt_creds = crt_s3_creds()
    except botocore.exceptions.NoCredentialsError:
        return False

    is_matching_identity = (
        boto3_creds.access_key == crt_creds.access_key_id
        and boto3_creds.secret_key == crt_creds.secret_access_key
        and boto3_creds.token == crt_creds.session_token
    )
    return is_matching_identity


def create_crt_transfer_manager(client, config):
    """Create a CRTTransferManager for optimized data transfer."""
    crt_s3_client = get_crt_s3_client(client, config)
    if is_crt_compatible_request(client, crt_s3_client):
        return CRTTransferManager(
            crt_s3_client.crt_client, BOTOCORE_CRT_SERIALIZER
        )
    return None