# Copyright 1999-2020 Alibaba Group Holding Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License 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.

import asyncio
import itertools
import logging
import time
from collections import namedtuple, OrderedDict
from dataclasses import dataclass
from typing import Dict, Optional, Tuple

from .... import oscar as mo
from .... import resource as mars_resource

logger = logging.getLogger(__name__)

QuotaDumpType = namedtuple('QuotaDumpType', 'allocations requests hold_sizes')


@dataclass
class QuotaRequest:
    req_size: Tuple
    delta: int
    req_time: float
    event: asyncio.Event


class QuotaActor(mo.Actor):
    @classmethod
    def gen_uid(cls, band):
        return f'{band}_quota'

    def __init__(self, total_size):
        super().__init__()
        self._requests = OrderedDict()

        self._total_size = total_size
        self._allocations = dict()
        self._total_allocated = 0

        self._hold_sizes = dict()
        self._total_hold = 0

    def _has_space(self, delta):
        return self._total_allocated + delta <= self._total_size

    def _log_allocate(self, msg, *args, **kwargs):
        args += (self._total_allocated, self._total_size)
        logger.debug(msg + ' Allocated: %s, Total size: %s', *args, **kwargs)

    async def request_batch_quota(self, batch):
        """
        Request for resources in a batch
        :param batch: the request dict in form {request_key: request_size, ...}
        :return: if request is returned immediately, return True, otherwise False
        """
        all_allocated = True
        # check if the request is already allocated
        for key, size in batch.items():
            if key not in self._allocations or size > self._allocations.get(key):
                all_allocated = False
                break

        self._log_allocate('Receive batch quota request %r on %s.', batch, self.uid)
        sorted_req = sorted(batch.items(), key=lambda tp: tp[0])
        keys = tuple(tp[0] for tp in sorted_req)
        quota_sizes = tuple(tp[1] for tp in sorted_req)
        delta = sum(v - self._allocations.get(k, 0) for k, v in batch.items())

        # if all requested and allocation can still be applied, apply directly
        if all_allocated and self._has_space(delta):
            self._log_allocate('Quota request %r already allocated.', batch)
            return

        if delta > self._total_size:
            raise ValueError(f'Cannot allocate quota size {delta} '
                             f'larger than total capacity {self._total_size}.')

        if keys in self._requests:
            event = self._requests[keys].event
        else:
            has_space = self._has_space(delta)
            if has_space and not self._requests:
                # if no previous requests, we can apply directly
                self._log_allocate('Quota request met for key %r on %s.', keys, self.uid)
                self.alter_allocations(keys, quota_sizes, allocate=True)
                return
            else:
                # current free space cannot satisfy the request, the request is queued
                if not has_space:
                    self._log_allocate('Quota request unmet for key %r on %s.', keys, self.uid)
                else:
                    self._log_allocate('Quota request queued for key %r on %s.', keys, self.uid)
                event = asyncio.Event()
                quota_request = QuotaRequest(quota_sizes, delta, time.time(), event)
                if keys not in self._requests:
                    self._requests[keys] = quota_request

        async def waiter():
            try:
                await event.wait()
            except asyncio.CancelledError as ex:
                await self.ref().remove_requests.tell(keys)
                raise ex
        return waiter()

    def remove_requests(self, keys):
        self._requests.pop(keys, None)
        self._process_requests()

    def hold_quotas(self, keys):
        """
        Mark request quota as already been hold
        :param keys: request keys
        """
        for key in keys:
            try:
                alloc_size = self._allocations[key]
            except KeyError:
                continue
            self._total_hold += alloc_size - self._hold_sizes.get(key, 0)
            self._hold_sizes[key] = alloc_size

    def release_quotas(self, keys):
        """
        Release allocated quota in batch
        :param keys: request keys
        """
        total_alloc_size = 0

        for key in keys:
            try:
                alloc_size = self._allocations.pop(key)
                total_alloc_size += alloc_size
            except KeyError:
                continue
            self._total_hold -= self._hold_sizes.pop(key, 0)

        self._total_allocated -= total_alloc_size
        if total_alloc_size:
            self._process_requests()
            self._log_allocate('Quota keys %s released on %s.', keys, self.uid)

    def dump_data(self):
        return QuotaDumpType(self._allocations, self._requests, self._hold_sizes)

    def get_allocated_size(self):
        # get total allocated size, for debug purpose
        return self._total_allocated

    def alter_allocations(self, keys, quota_sizes, handle_shrink=True,
                          allocate=False):
        """
        Alter multiple requests
        :param keys: keys to update
        :param quota_sizes: new quota sizes, if None, no changes will be made
        :param handle_shrink: if True and the quota size less than the original, process requests in the queue
        :param allocate: if True, will allocate resources for new items
        :return:
        """
        quota_sizes = quota_sizes or itertools.repeat(None)
        total_old_size, total_diff = 0, 0
        for k, s in zip(keys, quota_sizes):
            old_size = self._allocations.get(k, 0)
            size_diff = 0

            if not allocate and k not in self._allocations:
                total_old_size += old_size
                continue

            if s != old_size:
                s = int(s)
                size_diff = s - old_size
                self._total_allocated += size_diff
                self._allocations[k] = s
                try:
                    self._total_hold += s - self._hold_sizes[k]
                    self._hold_sizes[k] = s
                except KeyError:
                    pass

            total_old_size += old_size
            total_diff += size_diff
        if handle_shrink and total_diff < 0:
            self._process_requests()
        self._log_allocate('Quota keys %r applied on %s. Total old Size: %s, Total diff: %s,',
                           keys, self.uid, total_old_size, total_diff)

    def _process_requests(self):
        """
        Process quota requests in the queue
        """
        removed = []
        for k, req in self._requests.items():
            if self._has_space(req.delta):
                self.alter_allocations(k, req.req_size, handle_shrink=False, allocate=True)
                req.event.set()
                removed.append(k)
            else:
                # Quota left cannot satisfy the next request, we quit
                break
        for k in removed:
            self._requests.pop(k, None)


class MemQuotaActor(QuotaActor):
    """
    Actor handling worker memory quota
    """
    def __init__(self, soft_limit, hard_limit=None, refresh_time=None):
        super().__init__(soft_limit)
        self._hard_limit = hard_limit
        self._last_memory_available = 0
        self._refresh_time = refresh_time or 10

    async def __post_create__(self):
        self.ref().update_mem_stats.tell_delay(delay=self._refresh_time)

    def update_mem_stats(self):
        """
        Refresh memory usage
        """
        cur_mem_available = mars_resource.virtual_memory().available
        if cur_mem_available > self._last_memory_available:
            # memory usage reduced: try reallocate existing requests
            self._process_requests()
        self._last_memory_available = cur_mem_available
        self.ref().update_mem_stats.tell_delay(delay=self._refresh_time)

    def _has_space(self, delta):
        if self._hard_limit is None:
            return super()._has_space(delta)

        mem_stats = mars_resource.virtual_memory()
        # calc available physical memory
        available_size = mem_stats.available - max(0, mem_stats.total - self._hard_limit) \
            - (self._total_allocated - self._total_hold)
        if max(delta, 0) >= available_size:
            logger.warning('%s met hard memory limitation: request %d, available %d, hard limit %d',
                           self.uid, delta, available_size, self._hard_limit)
            return False
        return super()._has_space(delta)

    def _log_allocate(self, msg, *args, **kwargs):
        if self._hard_limit is None:
            return super()._log_allocate(msg, *args, **kwargs)

        mem_stats = mars_resource.virtual_memory()
        # calc available physical memory
        available_size = mem_stats.available - max(0, mem_stats.total - self._hard_limit) \
            - (self._total_allocated - self._total_hold)
        args += (self._total_allocated, self._total_size, mem_stats.available, available_size,
                 self._hard_limit, self._total_hold)

        logger.debug(
            msg + ' Allocated: %s, Total size: %s, Phy available: %s, Hard available: %s,'
                  ' Hard limit: %s, Holding: %s',
            *args, **kwargs
        )


class WorkerQuotaManagerActor(mo.Actor):
    def __init__(self, default_config: Dict, band_configs: Optional[Dict] = None):
        self._cluster_api = None
        self._default_config = default_config
        self._band_configs = band_configs or dict()

    async def __post_create__(self):
        from ...cluster.api import ClusterAPI
        self._cluster_api = await ClusterAPI.create(self.address)

        band_to_slots = await self._cluster_api.get_bands()
        for band in band_to_slots.keys():
            band_config = self._band_configs.get(band[1], self._default_config)
            await mo.create_actor(
                MemQuotaActor, **band_config,
                uid=MemQuotaActor.gen_uid(band[1]),
                address=self.address)
