# Copyright 2022-2023 XProbe Inc.
# derived from copyright 1999-2021 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 logging
import os
import time
from typing import Dict, List, NamedTuple, Set, Tuple

import psutil
import xoscar as mo
from xoscar.backends.allocate_strategy import IdleLabel
from xoscar.errors import NoFreeSlot, SlotStateError

from ....resource import Resource
from ....typing import BandType
from ...cluster import ClusterAPI, WorkerSlotInfo

logger = logging.getLogger(__name__)


class DispatchDumpType(NamedTuple):
    free_slots: Set
    fresh_slots: Set


class WorkerSlotManagerActor(mo.Actor):
    _band_slot_infos: Dict[str, List[WorkerSlotInfo]]

    def __init__(self):
        self._cluster_api = None
        self._global_resource_ref = None

        self._band_slot_managers = dict()  # type: Dict[str, mo.ActorRef]

    async def __post_create__(self):
        self._cluster_api = await ClusterAPI.create(self.address)

        band_to_resource = await self._cluster_api.get_bands()
        for band, resource in band_to_resource.items():
            self._band_slot_managers[band] = await mo.create_actor(
                BandSlotManagerActor,
                band,
                int(resource.num_cpus or resource.num_gpus),
                self._global_resource_ref,
                uid=BandSlotManagerActor.gen_uid(band[1]),
                address=self.address,
            )

    async def __pre_destroy__(self):
        await asyncio.gather(
            *[mo.destroy_actor(ref) for ref in self._band_slot_managers.values()]
        )


class BandSlotManagerActor(mo.Actor):
    _free_slots: Set[int]
    _fresh_slots: Set[int]

    @classmethod
    def gen_uid(cls, band_name: str):
        return f"{band_name}_band_slot_manager"

    def __init__(
        self, band: BandType, n_slots: int, global_resource_ref: mo.ActorRef = None
    ):
        super().__init__()
        self._cluster_api = None

        self._band = band
        self._band_name = band[1]
        self._global_resource_ref = global_resource_ref
        self._n_slots = n_slots

        self._semaphore = asyncio.Semaphore(0)
        self._slot_control_refs = dict()
        self._free_slots = set()
        self._fresh_slots = set()
        self._slot_kill_events = dict()

        self._restarting = False
        self._restart_done_event = asyncio.Event()

        self._session_stid_to_slot = dict()
        self._slot_to_session_stid = dict()
        self._last_report_time = time.time()

        self._slot_to_proc = dict()
        self._usage_upload_task = None

    async def __post_create__(self):
        try:
            self._cluster_api = await ClusterAPI.create(self.address)
        except mo.ActorNotExist:
            pass

        strategy = IdleLabel(self._band_name, f"worker_slot_control")
        for slot_id in range(self._n_slots):
            self._slot_control_refs[slot_id] = await mo.create_actor(
                BandSlotControlActor,
                self.ref(),
                self._band_name,
                slot_id,
                uid=BandSlotControlActor.gen_uid(self._band_name, slot_id),
                address=self.address,
                allocate_strategy=strategy,
            )
            self._fresh_slots.add(slot_id)

        self._upload_slot_usage_with_delay()

    async def __pre_destroy__(self):
        self._usage_upload_task.cancel()

    async def _get_global_resource_ref(self):
        if self._global_resource_ref is not None:
            return self._global_resource_ref

        from ..supervisor import GlobalResourceManagerActor

        try:
            [self._global_resource_ref] = await self._cluster_api.get_supervisor_refs(
                [GlobalResourceManagerActor.default_uid()]
            )
        except mo.ActorNotExist:
            self._global_resource_ref = None
        return self._global_resource_ref

    def get_slot_address(self, slot_id: int):
        return self._slot_control_refs[slot_id].address

    def get_subtask_slot(self, session_stid: Tuple[str, str]):
        return self._session_stid_to_slot.get(session_stid)

    async def acquire_free_slot(self, session_stid: Tuple[str, str], block=True):
        if not block and self._semaphore.locked():
            raise NoFreeSlot(f"No free slot for {session_stid}")
        yield self._semaphore.acquire()
        if self._restarting:
            yield self._restart_done_event.wait()

        slot_id = self._free_slots.pop()
        self._fresh_slots.difference_update([slot_id])
        self._slot_to_session_stid[slot_id] = session_stid
        self._session_stid_to_slot[session_stid] = slot_id
        logger.debug("Slot %d acquired for subtask %r", slot_id, session_stid)
        raise mo.Return(slot_id)

    def release_free_slot(self, slot_id: int, session_stid: Tuple[str, str]):
        acquired_session_stid = self._slot_to_session_stid.pop(slot_id, None)
        if acquired_session_stid is None:
            raise SlotStateError(f"Slot {slot_id} is not acquired.")
        if acquired_session_stid != session_stid:
            raise SlotStateError(
                f"Slot {slot_id} releasing state incorrect, "
                f"the acquired session_stid: {acquired_session_stid}, "
                f"the releasing session_stid: {session_stid}"
            )
        acquired_slot_id = self._session_stid_to_slot.pop(acquired_session_stid)
        assert (
            acquired_slot_id == slot_id
        ), f"{acquired_session_stid}: acquired_slot_id {acquired_slot_id} != slot_id {slot_id}"

        logger.debug("Slot %d released", slot_id)

        if slot_id not in self._free_slots:
            self._free_slots.add(slot_id)
            self._semaphore.release()

    def register_slot(self, slot_id: int, pid: int):
        try:
            self._fresh_slots.add(slot_id)
            if slot_id in self._slot_kill_events:
                event = self._slot_kill_events.pop(slot_id)
                event.set()
            if slot_id in self._slot_to_session_stid:
                # We should release the slot by one role, if the slot is
                # acquired by the SubtaskExecutionActor, then the slot
                # should be released by it, too.
                session_stid = self._slot_to_session_stid[slot_id]
                logger.info(
                    "Slot %s registered by pid %s, current acquired session_stid is %s",
                    slot_id,
                    pid,
                    session_stid,
                )
            else:
                if slot_id not in self._free_slots:
                    self._free_slots.add(slot_id)
                    self._semaphore.release()
        finally:
            # psutil may raises exceptions, but currently we can't handle the register exception,
            # so put it to the finally.
            # TODO(fyrestone): handle register_slot failure.
            self._slot_to_proc[slot_id] = proc = psutil.Process(pid)
            # collect initial stats for the process
            proc.cpu_percent(interval=None)

    async def _kill_slot(self, slot_id: int):
        if slot_id in self._slot_kill_events:
            await self._slot_kill_events[slot_id].wait()
            return

        event = self._slot_kill_events[slot_id] = asyncio.Event()
        # TODO(fyrestone): Make it more reliable. e.g. kill_actor
        # success but the actor does not restart.
        try:
            await mo.kill_actor(self._slot_control_refs[slot_id])
        except ConnectionError:
            pass
        await event.wait()

    async def kill_slot(self, slot_id: int):
        self._free_slots.difference_update([slot_id])
        yield self._kill_slot(slot_id)

    async def restart_free_slots(self):
        if self._restarting:
            yield self._restart_done_event.wait()
            return

        self._restart_done_event = asyncio.Event()
        self._restarting = True
        slot_ids = [
            slot_id for slot_id in self._free_slots if slot_id not in self._fresh_slots
        ]
        if slot_ids:
            yield asyncio.gather(*[self._kill_slot(slot_id) for slot_id in slot_ids])
            logger.info("%d idle slots restarted", len(slot_ids))

        self._restarting = False
        self._restart_done_event.set()

    def _upload_slot_usage_with_delay(self, delay: int = 1):
        self._usage_upload_task = self.ref().upload_slot_usages.tell_delay(
            periodical=True, delay=delay
        )

    async def upload_slot_usages(self, periodical: bool = False):
        delays = []
        slot_infos = []
        global_resource_ref = await self._get_global_resource_ref()

        if global_resource_ref is None:  # pragma: no cover
            if periodical:
                self._upload_slot_usage_with_delay()
            return

        for slot_id, proc in self._slot_to_proc.items():
            if slot_id not in self._slot_to_session_stid:
                continue
            session_id, subtask_id = self._slot_to_session_stid[slot_id]
            cpu_usage, gpu_usage, processor_usage = 0, 0, 0
            if self._band_name.startswith("gpu"):
                processor_usage = gpu_usage = 1
            else:
                try:
                    processor_usage = cpu_usage = (
                        proc.cpu_percent(interval=None) / 100.0
                    )
                except psutil.NoSuchProcess:  # pragma: no cover
                    continue
                except psutil.AccessDenied as e:  # pragma: no cover
                    logger.warning("Access denied when getting cpu percent: %s", e)
                    processor_usage = cpu_usage = 0.0

            slot_infos.append(
                WorkerSlotInfo(
                    slot_id=slot_id,
                    session_id=session_id,
                    subtask_id=subtask_id,
                    processor_usage=processor_usage,
                )
            )

            if global_resource_ref is not None:  # pragma: no branch
                # FIXME fix band slot mistake
                delays.append(
                    global_resource_ref.update_subtask_resources.delay(
                        self._band[1],
                        session_id,
                        subtask_id,
                        Resource(
                            num_cpus=max(1.0, cpu_usage), num_gpus=max(1.0, gpu_usage)
                        ),
                    )
                )

        if delays:  # pragma: no branch
            yield global_resource_ref.update_subtask_resources.batch(*delays)
        if self._cluster_api is not None:
            await self._cluster_api.set_band_slot_infos(self._band_name, slot_infos)

        if periodical:
            self._upload_slot_usage_with_delay()

    def dump_data(self):
        """
        Get all refs of slots of a queue
        """
        return DispatchDumpType(self._free_slots, self._fresh_slots)


class BandSlotControlActor(mo.Actor):
    @classmethod
    def gen_uid(cls, band_name: str, slot_id: int):
        return f"{band_name}_{slot_id}_band_slot_control"

    def __init__(self, manager_ref, band_name, slot_id: int):
        self._manager_ref = manager_ref
        self._band_name = band_name
        self._slot_id = slot_id
        self._report_task = None

    async def __post_create__(self):
        self._report_task = asyncio.create_task(self._report_slot_ready())

    async def _report_slot_ready(self):
        from ...cluster.api import ClusterAPI

        try:
            self._cluster_api = await ClusterAPI.create(self.address)
            await self._cluster_api.wait_node_ready()
        except mo.ActorNotExist:
            pass

        await mo.wait_actor_pool_recovered(self.address)
        await self._manager_ref.register_slot.tell(self._slot_id, os.getpid())
