"""Transfer from s3 to storage."""

import hashlib
import io
import json
import os

from boto3.s3 import transfer
from daemon_asset_copy import config
from daemon_asset_copy.connectors import logger, s3
from daemon_asset_copy.constants import constant as const
from daemon_asset_copy.constants import job_io_fields
from ddtrace import tracer
from PIL import Image


@tracer.wrap()
def transfer_from_s3_to_s3(inputs):
    """Transfer file from s3 to s3.

    Args:
        inputs (dict): Inputs from prior step

    Returns:
        dict: Adding new outputs

    """
    log = logger.get_current_logger()

    # source s3 client
    s3_source_bucket = inputs[job_io_fields.S3_SOURCE_BUCKET]
    s3_source_key = inputs[job_io_fields.S3_SOURCE_KEY]
    s3_source_client = s3.get_s3_client()

    # destination s3 client
    s3_destination_bucket = inputs[job_io_fields.S3_DESTINATION_BUCKET]
    s3_destination_key = inputs[job_io_fields.S3_DESTINATION_KEY]
    storage_creds = json.loads(os.getenv('STORAGE_CREDENTIALS'))
    s3_destination_client = s3.get_s3_client(storage_creds)
    s3_destination_config = transfer.TransferConfig(
        multipart_threshold=config.S3_UPLOAD_CHUNK_SIZE,  # default was 8Mb
        multipart_chunksize=config.S3_UPLOAD_CHUNK_SIZE,  # default was 8Mb
    )

    log.info(
        'Transferring file s3://{bucket}/{key}'
        ' to s3://{destination_bucket}/{destination_key}'.format(
            bucket=s3_source_bucket,
            key=s3_source_key,
            destination_bucket=s3_destination_bucket,
            destination_key=s3_destination_key,
        )
    )

    # copy from source to destination via streaming
    s3_source_response = s3_source_client.get_object(
        Bucket=s3_source_bucket, Key=s3_source_key
    )
    s3_destination_client.upload_fileobj(
        s3_source_response['Body'],
        s3_destination_bucket,
        s3_destination_key,
        Config=s3_destination_config,
    )

    # return extra metadata for asset details.
    s3_destination_file_name = s3_destination_key.split('/')[-1]
    s3_file_object = s3_destination_client.get_object(
        Bucket=s3_destination_bucket, Key=s3_destination_key
    )
    file_extension = s3_destination_key.split('.')[-1]
    log.info(f'Get metadata for file_extension: {file_extension}')
    return {
        **_get_file_metadata(file_extension, s3_file_object),
        job_io_fields.DESTINATION_FILE_NAME: s3_destination_file_name,
    }


def _get_file_metadata(file_extension, s3_file_object):
    """Read file and return md5 hash, image width, height, size etc.

    Note: Use this method if you are not reading the file for transfer only to get metadata.

    Args:
        file_extension (str): file extension.
        s3_file_object (dict): get_object result.

    Returns:
        dict: Adding new outputs

    """
    s3_file_size = s3_file_object['ContentLength']

    if file_extension == const.TIF:
        # non chunked read
        s3_streaming_file_object = io.BytesIO(s3_file_object['Body'].read())
        with Image.open(s3_streaming_file_object) as image:
            return {
                job_io_fields.FILE_SIZE: s3_file_size,
                job_io_fields.IMAGE_WIDTH: image.width,
                job_io_fields.IMAGE_HEIGHT: image.height,
                job_io_fields.MD5_HASH: hashlib.md5(image.tobytes()).hexdigest(),
            }

    if file_extension in const.METADATA_REQUIRED_ASSETS:
        hash_md5 = hashlib.md5()
        # chunked read
        for chunk in iter(
            lambda: s3_file_object['Body'].read(config.S3_DOWNLOAD_CHUNK_SIZE), b''
        ):
            hash_md5.update(chunk)

        return {
            job_io_fields.FILE_SIZE: s3_file_size,
            job_io_fields.MD5_HASH: hash_md5.hexdigest(),
        }

    return {}