import logging
import time
from typing import List, Optional, Set, Callable

from clarabot_nano.core.spawn.base import AbstractSpawn

NANO_INFO_LOGGER = logging.getLogger("nano/info")


class AbstractSpawner(object):
    def __init__(
        self,
        instance_number: int = 0,
        only_spawn_when_clean: bool = False,
        spawn_callback: Optional[Callable[[AbstractSpawn], None]] = None,
        dispose_callback: Optional[Callable[[AbstractSpawn], None]] = None,
    ):
        self.__instance_number: int = instance_number
        self.__instances: List[Optional[AbstractSpawn]] = [None] * self.__instance_number
        self.__disposing: Set[AbstractSpawn] = set()
        self.__only_spawn_when_clean = only_spawn_when_clean
        self.__spawn_callback = spawn_callback
        self.__dispose_callback = dispose_callback

    def set_instance_number(self, new_number):
        self.__instance_number = new_number
        self.__instances.extend([None] * (self.__instance_number - len(self.__instances)))

    def instances(self):
        return [i for i in self.__instances if i is not None]

    def manage(self, nano_config):
        """
        Create or dispose of instances as needed.
        """
        # dispose excess instances
        self._dispose_excess()
        # check current instances
        free_instances = []
        for i in range(self.__instance_number):
            if self.__instances[i] is None:
                if not self.__only_spawn_when_clean or not self.__disposing:
                    self.__spawn_wrapper(i, nano_config)
                else:
                    free_instances.append(i)
            elif not self.__instances[i].is_healthy():
                NANO_INFO_LOGGER.info("Disposing %s because of heartbeat failure", self.__instances[i].get_pid_prefix())
                if self.__dispose_callback:
                    self.__dispose_callback(self.__instances[i])
                self.__instances[i].dispose()
                self.__disposing.add(self.__instances[i])
                # spawn new at end or next cycle depending on disposed state & settings
                free_instances.append(i)
                self.__instances[i] = None
        # release disposed instances
        for instance in list(self.__disposing):
            if instance.wait(0):
                self.__disposing.remove(instance)
        # quick respawn if allowed
        if not self.__only_spawn_when_clean or not self.__disposing:
            for i in free_instances:
                self.__spawn_wrapper(i, nano_config)

    def _dispose_excess(self):
        while len(self.__instances) > self.__instance_number:
            instance = self.__instances.pop(-1)
            if instance:
                if self.__dispose_callback:
                    self.__dispose_callback(instance)
                instance.dispose()
                self.__disposing.add(instance)

    def cleanup(self, timeout=5.0):
        t = time.perf_counter() + timeout
        self.__instance_number = 0
        self._dispose_excess()
        for instance in list(self.__disposing):
            if instance.wait(max(0.0, t - time.perf_counter())):
                self.__disposing.remove(instance)
        return not self.__disposing

    def kill(self):
        self.__instance_number = 0
        self._dispose_excess()
        for instance in self.__disposing:
            instance.kill()

    def respawn(self):
        """
        Stop and discard current then start new spawns.
        """
        instance_number = self.__instance_number
        try:
            self.__instance_number = 0
            self._dispose_excess()
        finally:
            self.__instance_number = instance_number

    def __spawn_wrapper(self, index, nano_config):
        self.__instances[index] = self._spawn(index, nano_config)
        if self.__spawn_callback:
            self.__spawn_callback(self.__instances[index])

    def _spawn(self, index, nano_config):
        """
        Only use local configs from nano here. The remote configs shall be handled by the spawned processes as
        their changes will also be transmitted to them.
        """
        raise NotImplementedError()
