# 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
from typing import List, Dict, Union, Type, TypeVar

from .... import oscar as mo
from ....lib.aio import alru_cache
from ...core import NodeRole, BandType
from .core import AbstractClusterAPI

APIType = TypeVar('APIType', bound='ClusterAPI')
logger = logging.getLogger(__name__)


class ClusterAPI(AbstractClusterAPI):
    def __init__(self, address: str):
        self._address = address
        self._locator_ref = None
        self._uploader_ref = None
        self._node_info_ref = None

    async def _init(self):
        from ..locator import SupervisorLocatorActor
        from ..uploader import NodeInfoUploaderActor
        from ..supervisor.node_info import NodeInfoCollectorActor

        self._locator_ref = await mo.actor_ref(SupervisorLocatorActor.default_uid(),
                                               address=self._address)
        self._uploader_ref = await mo.actor_ref(NodeInfoUploaderActor.default_uid(),
                                                address=self._address)
        [self._node_info_ref] = await self.get_supervisor_refs(
            [NodeInfoCollectorActor.default_uid()])

    @classmethod
    @alru_cache(cache_exceptions=False)
    async def create(cls: Type[APIType], address: str) -> APIType:
        api_obj = cls(address)
        await api_obj._init()
        return api_obj

    async def get_supervisors(self, watch=False) -> List[str]:
        if watch:
            return await self._locator_ref.watch_supervisors()
        else:
            return await self._locator_ref.get_supervisors()

    async def get_supervisors_by_keys(self, keys: List[str], watch: bool = False) -> List[str]:
        """
        Get supervisor address hosting the specified key

        Parameters
        ----------
        keys
            key for a supervisor address
        watch
            if True, will watch changes of supervisor changes

        Returns
        -------
        out
            addresses of the supervisor
        """
        if not watch:
            get_supervisor = self._locator_ref.get_supervisor
            return await get_supervisor.batch(
                *(get_supervisor.delay(k) for k in keys)
            )
        else:
            return await self._locator_ref.watch_supervisors_by_keys(keys)

    async def get_supervisor_refs(self, uids: List[str], watch: bool = False) -> List[mo.ActorRef]:
        """
        Get actor references hosting the specified actor uid

        Parameters
        ----------
        uids
            uids for a supervisor address
        watch
            if True, will watch changes of supervisor changes

        Returns
        -------
        out : List[mo.ActorRef]
            references of the actors
        """
        addrs = await self.get_supervisors_by_keys(uids, watch=watch)
        return await asyncio.gather(*[
            mo.actor_ref(uid, address=addr) for addr, uid in zip(addrs, uids)
        ])

    async def watch_nodes(self, role: NodeRole, env: bool = False,
                          resource: bool = False, state: bool = False) -> List[Dict[str, Dict]]:
        return await self._node_info_ref.watch_nodes(
            role, env=env, resource=resource, state=state)

    async def get_nodes_info(self, nodes: List[str] = None, role: NodeRole = None,
                             env: bool = False, resource: bool = False, state: bool = False):
        return await self._node_info_ref.get_nodes_info(
            nodes=nodes, role=role, env=env, resource=resource, state=state)

    async def get_all_bands(self, role: NodeRole = None,
                            watch: bool = False) -> Dict[BandType, int]:
        if watch:
            return await self._node_info_ref.watch_all_bands(role)
        return await self._node_info_ref.get_all_bands(role)

    async def get_mars_versions(self) -> List[str]:
        return await self._node_info_ref.get_mars_versions()

    async def get_bands(self) -> Dict:
        """
        Get bands that can be used for computation on current node.

        Returns
        -------
        band_to_slots : dict
            Band to n_slot.
        """
        return await self._uploader_ref.get_bands()

    async def set_state_value(self, key: str, value: Union[List, Dict]):
        """
        Set state value for current node

        Parameters
        ----------
        key
            state key to set
        value
            value of the state
        """
        await self._uploader_ref.set_state_value.tell(key, value)

    async def set_band_resource(self, band: str, values: Dict):
        """
        Set resource usage for a band

        Parameters
        ----------
        band
        values

        Returns
        -------

        """
        await self._uploader_ref.set_band_resource.tell(band, values)

    async def mark_node_ready(self):
        """
        Mark current node ready for work loads
        """
        await self._uploader_ref.mark_node_ready()

    async def wait_all_supervisors_ready(self):
        """
        Wait till all expected supervisors are ready
        """
        await self._locator_ref.wait_all_supervisors_ready()


class MockClusterAPI(ClusterAPI):
    @classmethod
    async def create(cls: Type[APIType], address: str, **kw) -> APIType:
        from ..locator import SupervisorLocatorActor
        from ..uploader import NodeInfoUploaderActor
        from ..supervisor.node_info import NodeInfoCollectorActor

        dones, _ = await asyncio.wait([
            mo.create_actor(SupervisorLocatorActor, 'fixed', address,
                            uid=SupervisorLocatorActor.default_uid(),
                            address=address),
            mo.create_actor(NodeInfoCollectorActor,
                            uid=NodeInfoCollectorActor.default_uid(),
                            address=address),
            mo.create_actor(NodeInfoUploaderActor, NodeRole.WORKER,
                            interval=kw.get('upload_interval'),
                            band_to_slots=kw.get('band_to_slots'),
                            use_gpu=kw.get('use_gpu', False),
                            uid=NodeInfoUploaderActor.default_uid(),
                            address=address),
        ])

        for task in dones:
            try:
                task.result()
            except mo.ActorAlreadyExist:  # pragma: no cover
                pass

        api = await super().create(address=address)
        await api.mark_node_ready()
        return api
