import logging
from asyncio import gather
from asyncio.tasks import Task
from typing import Mapping, List, Iterable, Any

from .types import CoroutineFunc, FinalCallbackT, CancelCallbackT
from .task import start_task


log = logging.getLogger(__name__)


class TaskPool:
    def __init__(self, func: CoroutineFunc, args: Iterable[Any] = (), kwargs: Mapping[str, Any] = None,
                 final_callback: FinalCallbackT = None, cancel_callback: CancelCallbackT = None) -> None:
        self._func: CoroutineFunc = func
        self._args: Iterable[Any] = args
        self._kwargs: Mapping[str, Any] = kwargs if kwargs is not None else {}
        self._final_callback: FinalCallbackT = final_callback
        self._cancel_callback: CancelCallbackT = cancel_callback
        self._tasks: List[Task] = []

    @property
    def func_name(self) -> str:
        return self._func.__name__

    @property
    def size(self) -> int:
        return len(self._tasks)

    def __repr__(self) -> str:
        return f'<{self.__class__.__name__} func={self.func_name} size={self.size}>'

    def _task_name(self, i: int) -> str:
        return f'{self.func_name}_pool_task_{i}'

    def _start_one(self) -> None:
        self._tasks.append(start_task(self._func(*self._args, **self._kwargs), self._task_name(self.size),
                                      final_callback=self._final_callback, cancel_callback=self._cancel_callback))

    def start(self, num: int = 1) -> None:
        for _ in range(num):
            self._start_one()

    def stop(self, num: int = 1) -> int:
        if num < 1:
            return 0
        return sum(task.cancel() for task in reversed(self._tasks[-num:]))

    async def gather(self, return_exceptions: bool = False):
        results = await gather(*self._tasks, return_exceptions=return_exceptions)
        self._tasks = []
        return results