# 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 importlib
import logging
from typing import Dict, Optional, Type, Union

from .... import oscar as mo
from ...core import BandType
from ..core import Subtask, SubtaskResult
from ..errors import SlotOccupiedAlready
from .processor import SubtaskProcessor, SubtaskProcessorActor

logger = logging.getLogger(__name__)


SubtaskRunnerRef = Union["SubtaskRunnerActor", mo.ActorRef]


class SubtaskRunnerActor(mo.Actor):
    _session_id_to_processors: Dict[str, Union[mo.ActorRef, SubtaskProcessorActor]]
    _running_processor: Optional[Union[mo.ActorRef, SubtaskProcessorActor]]
    _last_processor: Optional[Union[mo.ActorRef, SubtaskProcessorActor]]

    @classmethod
    def gen_uid(cls, band_name: str, slot_id: int):
        return f'slot_{band_name}_{slot_id}_subtask_runner'

    def __init__(self,
                 supervisor_address: str,
                 band: BandType,
                 subtask_processor_cls: Type = None):
        self._supervisor_address = supervisor_address
        self._band = band
        self._subtask_processor_cls = \
            self._get_subtask_processor_cls(subtask_processor_cls)

        self._session_id_to_processors = dict()
        self._running_processor = None
        self._last_processor = None

    @classmethod
    def _get_subtask_processor_cls(cls, subtask_processor_cls):
        if subtask_processor_cls is None:
            return SubtaskProcessor
        else:
            assert isinstance(subtask_processor_cls, str)
            module, class_name = subtask_processor_cls.rsplit('.', 1)
            return getattr(importlib.import_module(module), class_name)

    async def _run_subtask(self, subtask: Subtask):
        processor = await self._init_subtask_processor(subtask)
        self._subtask_info.processor = processor
        return await processor.run()

    async def run_subtask(self, subtask: Subtask):
        if self._running_processor is not None:  # pragma: no cover
            subtask_id = await self._running_processor.get_running_subtask_id()
            # current subtask is still running
            raise SlotOccupiedAlready(
                f'There is subtask(id: {subtask_id}) running, '
                f'cannot run another subtask')

        session_id = subtask.session_id
        if session_id not in self._session_id_to_processors:
            self._session_id_to_processors[session_id] = await mo.create_actor(
                SubtaskProcessorActor, session_id, self._band,
                self._supervisor_address, self._subtask_processor_cls,
                uid=SubtaskProcessorActor.gen_uid(session_id),
                address=self.address)
        processor = self._session_id_to_processors[session_id]
        self._running_processor = self._last_processor = processor
        try:
            result = yield self._running_processor.run(subtask)
        finally:
            self._running_processor = None
        raise mo.Return(result)

    async def wait_subtask(self):
        yield self._running_processor.wait()

    async def get_subtask_result(self) -> SubtaskResult:
        return self._last_processor.result()

    def is_runner_free(self):
        return self._running_processor is None

    async def cancel_subtask(self):
        if self._running_processor is None:
            return
        yield self._running_processor.cancel()
