import concurrent
import hashlib
import logging
import math
from concurrent.futures.thread import ThreadPoolExecutor
from time import sleep

from botocore.exceptions import ClientError

from apollo_delphi_migration.exceptions import ValidationException, PostProcessingException
from apollo_delphi_migration.migration import MigrationTaskPostProcessor, MigrationContext
from apollo_delphi_migration.s3 import Bucket


class FileSizeValidationPostProcessor(MigrationTaskPostProcessor):
    __MAX_GET_SIZE_ATTEMPTS = 10

    def post_process(self, src_key: str, dst_key: str, context: MigrationContext):
        self._logger.info(
            '[POST_PROCESSING][SIZE_VALIDATION] Validation started (%s <-> %s). ',
            src_key,
            dst_key)

        # Get size of source file
        src_file_size = context.src_bucket.size(src_key)

        # Make several attempts to get destination file size.
        # Several attempts are required since S3 may unexpectedly return 404 for existing file.
        attempts = 0
        while True:
            try:
                dst_file_size = context.dst_bucket.size(dst_key)
                break
            except ClientError as exc:
                if exc.response['Error']['Code'] == '404' and attempts < self.__MAX_GET_SIZE_ATTEMPTS:
                    attempts += 1
                    self._logger.debug(
                        ('[POST_PROCESSING][SIZE_VALIDATION] '
                         'Destination file not found (%s). ClientError (404). Attempt %s'),
                        dst_key,
                        attempts)
                    sleep(0.5 * attempts)  # wait for N sec
                else:
                    raise

        if src_file_size != dst_file_size:
            raise ValidationException(f'File sizes are not equal ({src_file_size} != {dst_file_size})')

        self._logger.info('[POST_PROCESSING][SIZE_VALIDATION] Validation passed (%s == %s)', src_key, dst_key)


class AbstractFingerprintPostProcessor(MigrationTaskPostProcessor):

    def __init__(self,
                 logger: logging.Logger,
                 executor: ThreadPoolExecutor) -> None:
        super().__init__(logger)
        self._executor = executor

    def __hash_chunk(self, bucket: Bucket, key: str, chunk_index: int, chunk_size: int):
        read_from = chunk_size * chunk_index
        read_to = read_from + chunk_size

        self._logger.debug('[POST_PROCESSING] Reading bytes from %s. Start=%s. End=%s.', key, read_from, read_to)

        chunk = bucket.read(key, (read_from, read_to))
        md5 = hashlib.md5(chunk)
        digest = md5.digest()

        self._logger.debug(
            '[POST_PROCESSING] MD5 calculated for chunk (Key=%s, ChunkIndex=%s, ChunkSize=%s, MD5=%s)',
            key,
            chunk_index,
            len(chunk),
            md5.hexdigest())

        return chunk_index, digest, len(chunk)

    def _hash_object(self, bucket: Bucket, key: str, chunk_size: int):
        self._logger.debug('[POST_PROCESSING] Creating fingerprint for S3 object `%s`', key)

        object_size = bucket.size(key)

        futures = []
        for i in range(math.ceil(object_size / chunk_size)):
            futures.append(self._executor.submit(self.__hash_chunk, bucket, key, i, chunk_size))

        results = []
        total_hashed_bytes = 0
        for future in concurrent.futures.as_completed(futures):
            index, digest, hashed_bytes = future.result()
            total_hashed_bytes += hashed_bytes
            results.append((index, digest))

        if total_hashed_bytes != object_size:
            raise PostProcessingException(
                f'Number of hashed bytes ({total_hashed_bytes}) != s3 object size ({object_size})')

        md5 = hashlib.md5()
        for (index, digest) in sorted(results, key=lambda x: x[0]):
            self._logger.debug('[POST_PROCESSING] Aggregating chunk hashes (Key=%s, ChunkIndex=%s)', key, index)
            md5.update(digest)

        return md5.hexdigest()

    def post_process(self, src_key: str, dst_key: str, context):
        raise NotImplementedError


class FingerprintTaggingPostProcessor(AbstractFingerprintPostProcessor):
    DEFAULT_CHUNK_SIZE = 8 * 1024 * 1024  # 8MB

    def __init__(self,
                 logger: logging.Logger,
                 executor: ThreadPoolExecutor,
                 overwrite: bool = False,
                 chunk_size: int = DEFAULT_CHUNK_SIZE) -> None:
        super().__init__(logger, executor)
        self.__chunk_size = chunk_size
        self.__overwrite = overwrite

    def post_process(self, src_key: str, dst_key: str, context: MigrationContext):
        self._logger.info('[POST_PROCESSING][FINGERPRINT_TAGGING] Tagging procedure is started for %s', dst_key)

        if not self.__overwrite and context.dst_bucket.get_tag(dst_key, 'Fingerprint'):
            # Do nothing if we are not going to overwrite existing Fingerprint tag
            self._logger.info(('[POST_PROCESSING][FINGERPRINT_TAGGING] '
                               '`Fingerprint` tag already exists for %s key. Hash calculation is being skiped.'),
                              dst_key)
            return

        md5_hex = self._hash_object(context.dst_bucket, dst_key, self.__chunk_size)

        self._logger.debug(('[POST_PROCESSING][FINGERPRINT_TAGGING] '
                            'Fingerprint created (%s). Tag is being attached to S3 object %s'),
                           md5_hex,
                           dst_key)

        context.dst_bucket.add_tags(dst_key, Fingerprint=f'{md5_hex}.{self.__chunk_size}')

        self._logger.info(('[POST_PROCESSING][FINGERPRINT_TAGGING] '
                           '`Fingerprint` tag is attached to S3 object `%s`'),
                          dst_key)


class FingerprintValidationPostProcessor(AbstractFingerprintPostProcessor):

    def post_process(self, src_key: str, dst_key: str, context):
        self._logger.info(
            '[POST_PROCESSING][FINGERPRINT_VALIDATION] Validation started (%s <-> %s).',
            src_key,
            dst_key)

        fingerprint = context.dst_bucket.get_tag(dst_key, 'Fingerprint')
        if not fingerprint:
            raise ValidationException(f'The tag with key `Fingerprint` does not exist for dst s3 object')

        dst_md5_hex, chunk_size = tuple(fingerprint.split('.'))
        src_md5_hex = self._hash_object(context.src_bucket, src_key, int(chunk_size))

        if dst_md5_hex != src_md5_hex:
            raise ValidationException(f'Fingerprints are not equal (src: {src_md5_hex} != dst: {dst_md5_hex})')

        self._logger.info('[POST_PROCESSING][FINGERPRINT_VALIDATION] Passed: %s == %s', src_key, dst_key)