######################################################################
#
# File: b2sdk/transfer/emerge/planner/planner.py
#
# Copyright 2020 Backblaze Inc. All Rights Reserved.
#
# License https://www.backblaze.com/using_b2_code.html
#
######################################################################

import hashlib
import json

from abc import ABCMeta, abstractmethod
from collections import deque
from itertools import chain

from b2sdk.transfer.emerge.planner.part_definition import (
    CopyEmergePartDefinition,
    UploadEmergePartDefinition,
    UploadSubpartsEmergePartDefinition,
)
from b2sdk.transfer.emerge.planner.upload_subpart import (
    LocalSourceUploadSubpart,
    RemoteSourceUploadSubpart,
)

MEGABYTE = 1000 * 1000
GIGABYTE = 1000 * MEGABYTE


class UploadBuffer(object):
    """ data container used by EmergePlanner for temporary storage of write intents """

    def __init__(self, start_offset, buff=None):
        self._start_offset = start_offset
        self._buff = buff or []
        if self._buff:
            self._end_offset = self._buff[-1][1]
        else:
            self._end_offset = self._start_offset

    @property
    def start_offset(self):
        return self._start_offset

    @property
    def end_offset(self):
        return self._end_offset

    @property
    def length(self):
        return self.end_offset - self.start_offset

    def intent_count(self):
        return len(self._buff)

    def get_intent(self, index):
        return self.get_item(index)[0]

    def get_item(self, index):
        return self._buff[index]

    def iter_items(self):
        return iter(self._buff)

    def append(self, intent, fragment_end):
        self._buff.append((intent, fragment_end))
        self._end_offset = fragment_end

    def get_slice(self, start_idx=None, end_idx=None, start_offset=None):
        start_idx = start_idx or 0
        buff_slice = self._buff[start_idx:end_idx]
        if start_offset is None:
            if start_idx == 0:
                start_offset = self.start_offset
            else:
                start_offset = self._buff[start_idx - 1:start_idx][0][1]
        return self.__class__(start_offset, buff_slice)


class EmergePlanner(object):
    """ Creates a list of actions required for advanced creation of an object in the cloud from an iterator of write intent objects """
    DEFAULT_MIN_PART_SIZE = 5 * MEGABYTE
    DEFAULT_RECOMMENDED_UPLOAD_PART_SIZE = 100 * MEGABYTE
    DEFAULT_MAX_PART_SIZE = 5 * GIGABYTE

    def __init__(
        self,
        min_part_size=None,
        recommended_upload_part_size=None,
        max_part_size=None,
    ):
        self.min_part_size = min_part_size or self.DEFAULT_MIN_PART_SIZE
        self.recommended_upload_part_size = recommended_upload_part_size or self.DEFAULT_RECOMMENDED_UPLOAD_PART_SIZE
        self.max_part_size = max_part_size or self.DEFAULT_MAX_PART_SIZE
        assert self.min_part_size <= self.recommended_upload_part_size <= self.max_part_size

    @classmethod
    def from_account_info(
        cls,
        account_info,
        min_part_size=None,
        recommended_upload_part_size=None,
        max_part_size=None
    ):
        # TODO: add support for getting `min_part_size` and `max_part_size` from account info
        if recommended_upload_part_size is None:
            # TODO: change `get_minimum_part_size` to correct name
            recommended_upload_part_size = account_info.get_minimum_part_size()
        if min_part_size is None and recommended_upload_part_size < cls.DEFAULT_MIN_PART_SIZE:
            min_part_size = recommended_upload_part_size
        if max_part_size is None and recommended_upload_part_size > cls.DEFAULT_MAX_PART_SIZE:
            max_part_size = recommended_upload_part_size
        kwargs = {
            'min_part_size': min_part_size,
            'recommended_upload_part_size': recommended_upload_part_size,
            'max_part_size': max_part_size,
        }
        return cls(**{key: value for key, value in kwargs.items() if value is not None})

    def get_emerge_plan(self, write_intents):
        write_intents = sorted(write_intents, key=lambda intent: intent.destination_offset)
        return self._get_emerge_plan(write_intents, EmergePlan)

    def get_streaming_emerge_plan(self, write_intent_iterator):
        return self._get_emerge_plan(write_intent_iterator, StreamingEmergePlan)

    def _get_emerge_plan(self, write_intent_iterator, plan_class):
        return plan_class(
            self._get_emerge_parts(
                self._select_intent_fragments(self._validatation_iterator(write_intent_iterator))
            )
        )

    def _get_emerge_parts(self, intent_fragments_iterator):
        # This is where the magic happens. Instead of describing typical inputs and outputs here,
        # We've put them in tests. It is recommended to read those tests before trying to comprehend
        # the implementation details of this function.
        min_part_size = self.min_part_size

        # this stores current intent that we need to process - we may get
        # it in fragments so we want to glue just by updating `current_end`
        current_intent = None
        current_end = 0

        upload_buffer = UploadBuffer(0)
        first = True
        last = False
        for intent, fragment_end in intent_fragments_iterator:
            if current_intent is None:
                # this is a first loop run - just initialize current intent
                current_intent = intent
                current_end = fragment_end
                continue

            if intent is current_intent:
                # new intent is the same as previously processed intent, so lets glue them together
                # this happens when the caller splits N overlapping intents into overlapping fragments
                # and two fragments from the same intent end up streaming into here one after the other
                current_end = fragment_end
                continue

            if intent is None:
                last = True

            # incoming intent is different - this means that now we have to decide what to do:
            # if this is a copy intent and we want to copy it server-side, then we have to
            # flush the whole upload buffer we accumulated so far, but OTOH we may decide that we just want to
            # append it to upload buffer (see complete, untrivial logic below) and then maybe
            # flush some upload parts from upload bufffer (if there is enough in the buffer)

            current_len = current_end - upload_buffer.end_offset
            # should we flush the upload buffer or do we have to add a chunk of the copy first?
            if current_intent.is_copy() and current_len >= min_part_size:
                # check if we can flush upload buffer or there is some missing bytes to fill it to `min_part_size`
                if upload_buffer.intent_count() > 0 and upload_buffer.length < min_part_size:
                    missing_length = min_part_size - upload_buffer.length
                else:
                    missing_length = 0
                if missing_length > 0 and current_len - missing_length < min_part_size:
                    # current intent is *not* a "small copy", but upload buffer is small
                    # and current intent is too short with the buffer to reach the minimum part size
                    # so we append current intent to upload buffer
                    upload_buffer.append(current_intent, current_end)
                else:
                    if missing_length > 0:
                        # we "borrow" a fragment of current intent to upload buffer
                        # to fill it to minimum part size
                        upload_buffer.append(
                            current_intent, upload_buffer.end_offset + missing_length
                        )
                    # completely flush the upload buffer
                    for upload_buffer_part in self._buff_split(upload_buffer):
                        yield self._get_upload_part(upload_buffer_part)
                    # split current intent (copy source) to parts and yield
                    copy_parts = self._get_copy_parts(
                        current_intent,
                        start_offset=upload_buffer.end_offset,
                        end_offset=current_end,
                    )
                    for part in copy_parts:
                        yield part
                    upload_buffer = UploadBuffer(current_end)
            else:
                if current_intent.is_copy() and first and last:
                    # this is a single intent "small copy" - we force use of `copy_file`
                    copy_parts = self._get_copy_parts(
                        current_intent,
                        start_offset=upload_buffer.end_offset,
                        end_offset=current_end,
                    )
                    for part in copy_parts:
                        yield part
                else:
                    # this is a upload source or "small copy" source (that is *not* single intent)
                    # either way we just add them to upload buffer
                    upload_buffer.append(current_intent, current_end)
                    upload_buffer_parts = list(self._buff_split(upload_buffer))
                    # we flush all parts excluding last one - we may want to extend
                    # this last part with "incoming" intent in next loop run
                    for upload_buffer_part in upload_buffer_parts[:-1]:
                        yield self._get_upload_part(upload_buffer_part)
                    upload_buffer = upload_buffer_parts[-1]
            current_intent = intent
            first = False
            current_end = fragment_end
            if current_intent is None:
                # this is a sentinel - there would be no more fragments - we have to flush upload buffer
                for upload_buffer_part in self._buff_split(upload_buffer):
                    yield self._get_upload_part(upload_buffer_part)

    def _get_upload_part(self, upload_buffer):
        """ Build emerge part from upload buffer. """
        if upload_buffer.intent_count() == 1 and upload_buffer.get_intent(0).is_upload():
            intent = upload_buffer.get_intent(0)
            relative_offset = upload_buffer.start_offset - intent.destination_offset
            length = upload_buffer.length
            definition = UploadEmergePartDefinition(intent.outbound_source, relative_offset, length)
        else:
            subparts = []
            fragment_start = upload_buffer.start_offset
            for intent, fragment_end in upload_buffer.iter_items():
                relative_offset = fragment_start - intent.destination_offset
                length = fragment_end - fragment_start
                if intent.is_upload():
                    subpart_class = LocalSourceUploadSubpart
                elif intent.is_copy():
                    subpart_class = RemoteSourceUploadSubpart
                else:
                    raise RuntimeError('This cannot happen!!!')
                subparts.append(subpart_class(intent.outbound_source, relative_offset, length))
                fragment_start = fragment_end
            definition = UploadSubpartsEmergePartDefinition(subparts)
        return EmergePart(definition)

    def _get_copy_parts(self, copy_intent, start_offset, end_offset):
        """ Split copy intent to emerge parts. """
        fragment_length = end_offset - start_offset
        part_count = int(fragment_length / self.max_part_size)
        last_part_length = fragment_length % self.max_part_size
        if last_part_length == 0:
            last_part_length = self.max_part_size
        else:
            part_count += 1

        if part_count == 1:
            part_sizes = [last_part_length]
        else:
            if last_part_length < int(fragment_length / (part_count + 1)):
                part_count += 1
            base_part_size = int(fragment_length / part_count)
            size_remainder = fragment_length % part_count
            part_sizes = [
                base_part_size + (1 if i < size_remainder else 0) for i in range(part_count)
            ]

        copy_source = copy_intent.outbound_source
        relative_offset = start_offset - copy_intent.destination_offset
        for part_size in part_sizes:
            yield EmergePart(CopyEmergePartDefinition(copy_source, relative_offset, part_size))
            relative_offset += part_size

    def _buff_split(self, upload_buffer):
        """ Split upload buffer to parts candidates - smaller upload buffers.

        :rtype iterator[b2sdk.transfer.emerge.planner.planner.UploadBuffer]:
        """
        if upload_buffer.intent_count() == 0:
            return
        tail_buffer = upload_buffer
        while True:
            if tail_buffer.length < self.recommended_upload_part_size + self.min_part_size:
                # `EmergePlanner_buff_partition` can split in such way that tail part
                # can be smaller than `min_part_size` - to avoid unnecessary download of possible
                # incoming copy intent, we don't split futher
                yield tail_buffer
                return
            head_buff, tail_buffer = self._buff_partition(tail_buffer)
            yield head_buff

    def _buff_partition(self, upload_buffer):
        """ Split upload buffer to two parts (smaller upload buffers).

        In result left part cannot be splitted more, and nothing can be assumed about right part.

        :rtype tuple(b2sdk.transfer.emerge.planner.planner.UploadBuffer,
                     b2sdk.transfer.emerge.planner.planner.UploadBuffer):
        """
        left_buff = UploadBuffer(upload_buffer.start_offset)
        buff_start = upload_buffer.start_offset
        for idx, (intent, fragment_end) in enumerate(upload_buffer.iter_items()):
            candidate_size = fragment_end - buff_start
            if candidate_size > self.recommended_upload_part_size:
                right_fragment_size = candidate_size - self.recommended_upload_part_size
                left_buff.append(intent, fragment_end - right_fragment_size)
                return left_buff, upload_buffer.get_slice(
                    start_idx=idx, start_offset=left_buff.end_offset
                )
            else:
                left_buff.append(intent, fragment_end)
                if candidate_size == self.recommended_upload_part_size:
                    return left_buff, upload_buffer.get_slice(start_idx=idx + 1)

        return left_buff, UploadBuffer(left_buff.end_offset)

    def _select_intent_fragments(self, write_intent_iterator):
        """ Select overapping write intent fragments to use.

        To solve overlapping intents selection, intents can be split to smaller fragments.
        Those fragments are yieled as soon as decision can be made to use them,
        so there is possibility that one intent is yielded in multiple fragments. Those
        would be merged again by higher level iterator that produces emerge parts, but
        in principle this merging can happen here. Not merging it is a code design decision
        to make this function easier to implement and also it would allow yielding emerge parts
        a bit quickier.
        """

        # `protected_intent_length` for upload state is 0, so it would generate at most single intent fragment
        # every loop iteration, but algorithm is not assuming that - one may one day choose to
        # protect upload fragments length too - eg. to avoid situation when file is opened to
        # read just small number of bytes and then switch to another overlapping upload source
        upload_intents_state = IntentsState()
        copy_intents_state = IntentsState(protected_intent_length=self.min_part_size)

        last_sent_offset = 0
        incoming_offset = None
        while True:
            incoming_intent = next(write_intent_iterator, None)
            if incoming_intent is None:
                incoming_offset = None
            else:
                incoming_offset = incoming_intent.destination_offset

            upload_intents = list(
                upload_intents_state.state_update(last_sent_offset, incoming_offset)
            )
            copy_intents = list(copy_intents_state.state_update(last_sent_offset, incoming_offset))

            intent_fragments = self._merge_intent_fragments(
                last_sent_offset,
                upload_intents,
                copy_intents,
            )

            for intent, intent_fragment_end in intent_fragments:
                yield intent, intent_fragment_end
                last_sent_offset = intent_fragment_end

            if incoming_offset is not None and last_sent_offset < incoming_offset:
                raise ValueError(
                    'Cannot emerge file with holes. '
                    'Found hole range: ({}, {})'.format(last_sent_offset, incoming_offset)
                )

            if incoming_intent is None:
                yield None, None  # lets yield sentinel for cleaner `_get_emerge_parts` implementation
                return
            if incoming_intent.is_upload():
                upload_intents_state.add(incoming_intent)
            elif incoming_intent.is_copy():
                copy_intents_state.add(incoming_intent)
            else:
                raise RuntimeError('This should not happen at all!')

    def _merge_intent_fragments(self, start_offset, upload_intents, copy_intents):
        """ Select "competing" upload and copy fragments.

        Upload and copy fragments may overlap so we nedd to choose right one
        to use - copy fragments are prioritized unless this fragment is unprotected
        (we use "protection" as an abstract for "short copy" fragments - meaning upload
        fragments have higher priority than "short copy")
        """
        upload_intents = deque(upload_intents)
        copy_intents = deque(copy_intents)
        while True:
            upload_intent = copy_intent = None
            if upload_intents:
                upload_intent, upload_end, _ = upload_intents[0]
            if copy_intents:
                copy_intent, copy_end, copy_protected = copy_intents[0]

            if upload_intent is not None and copy_intent is not None:
                if not copy_protected:
                    yield_intent = upload_intent
                else:
                    yield_intent = copy_intent
                start_offset = min(upload_end, copy_end)
                yield yield_intent, start_offset
                if start_offset >= upload_end:
                    upload_intents.popleft()
                if start_offset >= copy_end:
                    copy_intents.popleft()
            elif upload_intent is not None:
                yield upload_intent, upload_end
                upload_intents.popleft()
            elif copy_intent is not None:
                yield copy_intent, copy_end
                copy_intents.popleft()
            else:
                return

    def _validatation_iterator(self, write_intents):
        """ Iterate over write intents and validate length and order. """
        last_offset = 0
        for write_intent in write_intents:
            if write_intent.length is None:
                raise ValueError('Planner cannot support write intents of unknown length')
            if write_intent.destination_offset < last_offset:
                raise ValueError('Write intent stream have to be sorted by destination offset')
            last_offset = write_intent.destination_offset
            yield write_intent


class IntentsState(object):
    """ Store and process state of incoming write intents to solve
    overlapping intents selection in streaming manner.

    It does not check if intents are of the same kind (upload/copy), but the intention
    was to use it to split incoming intents by kind (two intents state are required then).
    If there would be no need for differentiating incoming intents, this would
    still work - so intent kind is ignored at this level. To address "short copy"
    prioritization problem (and avoidance) - ``protected_intent_length`` param was introduced
    to prevent logic from allowing too small fragments (if it is possible)
    """

    def __init__(self, protected_intent_length=0):
        self.protected_intent_length = protected_intent_length
        self._current_intent = None
        self._next_intent = None
        self._last_sent_offset = 0
        self._incoming_offset = None
        self._current_intent_start = None
        self._current_intent_end = None
        self._next_intent_end = None

    def add(self, incoming_intent):
        """ Add incoming intent to state.

        It has to called *after* ``IntentsState.state_update`` but it is not verified.
        """
        if self._next_intent is None:
            self._set_next_intent(incoming_intent)
        elif incoming_intent.destination_end_offset > self._next_intent_end:
            # here either incoming intent starts at the same position as next intent
            # (and current intent is None in such case - it was just cleared in `state_update`
            # or it was cleared some time ago - in previous iteratios) or we are in situation
            # when current and next intent overlaps, and `last_sent_offset` is now set to
            # incoming intent `destination_offset` - in both cases we want to choose
            # intent which has larger `destination_end_offset`
            self._set_next_intent(incoming_intent)

    def state_update(self, last_sent_offset, incoming_offset):
        """ Update the state using incoming intent offset.

        It has to be called *before* ``IntentsState.add`` and even if incoming intent
        would not be added to this intents state. It would yield a state of this stream
        of intents (like copy or upload) from ``last_sent_offset`` to ``incoming_offset``.
        So here happens the first stage of solving overlapping intents selection - but
        write intent iterator can be splitted to multiple substreams (like copy and upload)
        so additional stage is required to cover this.
        """
        if self._current_intent is not None:
            if last_sent_offset >= self._current_intent_end:
                self._set_current_intent(None, None)

        # `effective_incoming_offset` is a safeguard after intent iterator is drained
        if incoming_offset is not None:
            effective_incoming_offset = incoming_offset
        elif self._next_intent is not None:
            effective_incoming_offset = self._next_intent_end
        elif self._current_intent is not None:
            effective_incoming_offset = self._current_intent_end
        else:
            # intent iterator is drained and this state is empty
            return

        if (
            self._current_intent is None and self._next_intent is not None and (
                self._next_intent.destination_offset != effective_incoming_offset or
                incoming_offset is None
            )
        ):
            self._set_current_intent(self._next_intent, last_sent_offset)
            self._set_next_intent(None)

        # current and next can be both not None at this point only if they overlap
        if (
            self._current_intent is not None and self._next_intent is not None and
            effective_incoming_offset > self._current_intent_end
        ):
            # incoming intent does not overlap with current intent
            # so we switch to next because we are sure that we will have to use it anyway
            # (of course other overriding (eg. "copy" over "upload") state can have
            # overlapping intent but we have no information about it here)
            # but we also need to protect current intent length
            if not self._is_current_intent_protected():
                # we were unable to protect current intent, so we can safely rotate
                self._set_current_intent(self._next_intent, last_sent_offset)
                self._set_next_intent(None)
            else:
                remaining_len = self.protected_intent_length - (
                    last_sent_offset - self._current_intent_start
                )
                if remaining_len > 0:
                    last_sent_offset += remaining_len
                    if not self._can_be_protected(last_sent_offset, self._next_intent_end):
                        last_sent_offset = self._current_intent_end
                    yield self._current_intent, last_sent_offset, True
                self._set_current_intent(self._next_intent, last_sent_offset)
                self._set_next_intent(None)

        if self._current_intent is not None:
            yield (
                self._current_intent,
                min(effective_incoming_offset, self._current_intent_end),
                self._is_current_intent_protected(),
            )

    def _set_current_intent(self, intent, start_offset):
        self._current_intent = intent
        if self._current_intent is not None:
            self._current_intent_end = self._current_intent.destination_end_offset
        else:
            self._current_intent_end = None
            assert start_offset is None
        self._current_intent_start = start_offset

    def _set_next_intent(self, intent):
        self._next_intent = intent
        if self._next_intent is not None:
            self._next_intent_end = self._next_intent.destination_end_offset
        else:
            self._next_intent_end = None

    def _is_current_intent_protected(self):
        """ States if current intent is protected.

        Intent can be split to smaller fragments, but to choose upload over "small copy"
        we need to know for fragment if it is a "small copy" or not. In result of solving
        overlapping intents selection there might be a situation when original intent was not
        a small copy, but in effect it will be used only partially and in effect it may be a "small copy".
        Algorithm attempts to aviod using smaller fragments than ``protected_intent_length`` but
        sometimes it may be impossible. So if this function returns ``False`` it means
        that used length of this intent is smaller than ``protected_intent_length`` and the algorithm
        was unable to avoid this.
        """
        return self._can_be_protected(self._current_intent_start, self._current_intent_end)

    def _can_be_protected(self, start, end):
        return end - start >= self.protected_intent_length


class BaseEmergePlan(metaclass=ABCMeta):
    def __init__(self, emerge_parts):
        self.emerge_parts = emerge_parts

    @abstractmethod
    def is_large_file(self):
        pass

    @abstractmethod
    def get_total_length(self):
        pass

    @abstractmethod
    def get_plan_id(self):
        pass

    def enumerate_emerge_parts(self):
        return enumerate(self.emerge_parts, 1)


class EmergePlan(BaseEmergePlan):
    def __init__(self, emerge_parts):
        super(EmergePlan, self).__init__(list(emerge_parts))
        self._is_large_file = len(self.emerge_parts) > 1

    def is_large_file(self):
        return self._is_large_file

    def get_total_length(self):
        return sum(emerge_part.get_length() for emerge_part in self.emerge_parts)

    def get_plan_id(self):
        if all(part.is_hashable() for part in self.emerge_parts):
            return None

        json_id = json.dumps([emerge_part.get_part_id() for emerge_part in self.emerge_parts])
        return hashlib.sha1(json_id.encode()).hexdigest()


class StreamingEmergePlan(BaseEmergePlan):
    def __init__(self, emerge_parts_iterator):
        emerge_parts, self._is_large_file = self._peek_for_large_file(emerge_parts_iterator)
        super(StreamingEmergePlan, self).__init__(emerge_parts)

    def is_large_file(self):
        return self._is_large_file

    def get_total_length(self):
        return None

    def get_plan_id(self):
        return None

    def _peek_for_large_file(self, emerge_parts_iterator):
        first_part = next(emerge_parts_iterator, None)
        if first_part is None:
            raise ValueError('Empty emerge parts iterator')

        second_part = next(emerge_parts_iterator, None)
        if second_part is None:
            return iter([first_part]), False
        else:
            return chain([first_part, second_part], emerge_parts_iterator), True


class EmergePart(object):
    def __init__(self, part_definition, verification_ranges=None):
        self.part_definition = part_definition
        self.verification_ranges = verification_ranges

    def __repr__(self):
        return '<{classname} part_definition={part_definition}>'.format(
            classname=self.__class__.__name__,
            part_definition=repr(self.part_definition),
        )

    def get_length(self):
        return self.part_definition.get_length()

    def get_execution_step(self, execution_step_factory):
        return self.part_definition.get_execution_step(execution_step_factory)

    def get_part_id(self):
        return self.part_definition.get_part_id()

    def is_hashable(self):
        return self.part_definition.is_hashable()

    def get_sha1(self):
        return self.part_definition.get_sha1()
