###############################################################################
# This file is part of the tango-takeoff project.
#
# Copyright Lund University
#
# Distributed under the GNU GPLv3 license. See LICENSE file for more info.
###############################################################################
import asyncio
from typing import Callable

from tango.asyncio_executor import set_global_executor

from takeoff.custom_types import Host, Level, Order
from takeoff.error import TakeOffError

# Timeout (in seconds) to prevent waiting forever on asyncio tasks
# This is only as last resort
# We don't want this to be triggered as it's difficult to know what happened
# Errors should be caught in the task itself instead
# Value SHALL be greater than all the timouts and retries in the tasks
TASK_TIMEOUT = 120


class Worker:
    """
    A worker that parallelises hosts and levels.

    The worker encapsulate the asyncio work that groups and parallelises based
    on hosts and levels. Exactly what kind of work is going to be done is left
    to the callback, that way the grouping and parallelism can be re-used by
    different parts of the system.
    """

    def work(self, hosts: list[Host], order: Order, callback: Callable):
        asyncio.run(self._work(hosts, order, callback))
        # Make pytango forget about the loop
        set_global_executor(None)

    async def _work(self, hosts: list[Host], order: Order, callback: Callable):
        tasks = []
        for host in hosts:
            levels = sorted(host.levels, key=lambda x: x.level)
            if order == Order.DESCENDING:
                levels.reverse()

            tasks.append(self._work_on_levels(host.hostname, levels, callback))

        await asyncio.gather(*tasks)

    async def _work_on_levels(
        self, hostname: str, levels: list[Level], callback: Callable
    ):
        for level, servers in levels:
            tasks = []
            for name in servers:
                tasks.append(callback(hostname, level, name))
            try:
                await asyncio.wait_for(asyncio.gather(*tasks), timeout=TASK_TIMEOUT)
            except asyncio.TimeoutError:
                raise TakeOffError(
                    f"Task(s) on {hostname} took too long! Cancelled..."
                ) from None
