#
# This file is part of the Ingram Micro CloudBlue Connect EaaS Extension Runner.
#
# Copyright (c) 2021 Ingram Micro. All Rights Reserved.
#
import asyncio
import logging


from connect.eaas.core.proto import (
    Message,
    MessageType,
    SetupRequest,
)
from connect.eaas.runner.workers.base import WorkerBase
from connect.eaas.runner.constants import (
    DELAY_ON_CONNECT_EXCEPTION_SECONDS,
    RESULT_SENDER_MAX_RETRIES,
    RESULT_SENDER_WAIT_GRACE_SECONDS,
)
from connect.eaas.runner.helpers import get_version
from connect.eaas.runner.managers import (
    BackgroundTasksManager,
    InteractiveTasksManager,
    ScheduledTasksManager,
)


logger = logging.getLogger(__name__)


class Worker(WorkerBase):
    """
    The Worker is responsible to handle the websocket connection
    with the server. It will send the extension capabilities to
    the server and wait for tasks that need to be processed using
    the tasks manager.
    """
    def __init__(self, config, handler, runner_type=None):
        super().__init__(config)
        self.runner_type = runner_type
        self.handler = handler
        self.results_queue = asyncio.Queue()
        self.background_manager = BackgroundTasksManager(
            self.config,
            self.handler,
            self.results_queue.put,
        )
        self.interactive_manager = InteractiveTasksManager(
            self.config,
            self.handler,
            self.results_queue.put,
        )
        self.scheduled_manager = ScheduledTasksManager(
            self.config,
            self.handler,
            self.results_queue.put,
        )
        self.ws = None
        self.main_task = None
        self.results_task = None

    @property
    def running_tasks(self):
        return (
            self.background_manager.running_tasks
            + self.interactive_manager.running_tasks
            + self.scheduled_manager.running_tasks
        )

    def get_url(self):
        url = self.config.get_events_ws_url()
        url = f'{url}?running_tasks={self.background_manager.running_tasks}'
        url = f'{url}&running_scheduled_tasks={self.scheduled_manager.running_tasks}'
        if self.runner_type:
            url = f'{url}&runner_type={self.runner_type}'
        return url

    def get_setup_request(self):
        return Message(
            version=2,
            message_type=MessageType.SETUP_REQUEST,
            data=SetupRequest(
                event_subscriptions={
                    event_type: event['statuses'] or []
                    for event_type, event in self.handler.events.items()
                },
                variables=self.handler.variables,
                schedulables=self.handler.schedulables,
                repository={
                    'readme_url': self.handler.readme,
                    'changelog_url': self.handler.changelog,
                },
                runner_version=get_version(),
            ),
        ).dict()

    async def process_message(self, data):
        """
        Process a message received from the websocket server.
        """
        message = Message.deserialize(data)
        if message.message_type == MessageType.SETUP_RESPONSE:
            self.process_setup_response(message.data)
        elif message.message_type == MessageType.TASK:
            await self.process_task(message.data)
        elif message.message_type == MessageType.SHUTDOWN:
            await self.shutdown()

    async def process_task(self, task_data):
        """Send a task to a manager based on task category."""
        logger.info(
            f'received new {task_data.options.task_category} task: {task_data.options.task_id}',
        )
        manager = getattr(self, f'{task_data.options.task_category}_manager')
        await manager.submit(task_data)
        logger.info(f'task {task_data.options.task_id} submitted for processing')

    async def result_sender(self):  # noqa: CCR001
        """
        Dequeues results from the results queue and send it to
        the EaaS backend.
        """
        await self.run_event.wait()
        while True:
            if self.results_queue.empty():
                if not self.run_event.is_set() and self.running_tasks == 0:
                    logger.info('Worker exiting and no more running tasks: exit!')
                    return
                await asyncio.sleep(.5)
                continue
            logger.info(
                f'Current processing status: running={self.running_tasks} '
                f'results={self.results_queue.qsize()}',
            )
            result = await self.results_queue.get()
            logger.info(f'Got a result from queue: {result.options.task_id}')
            retries = 0
            while retries < RESULT_SENDER_MAX_RETRIES:
                try:
                    message = Message(
                        version=2,
                        message_type=MessageType.TASK,
                        data=result,
                    )
                    await self.ensure_connection()
                    await self.send(message.serialize())
                    logger.info(f'Result for task {result.options.task_id} has been sent.')
                    break
                except Exception:
                    logger.warning(
                        f'Attempt {retries} to send results for task '
                        f'{result.options.task_id} has failed.',
                    )
                    retries += 1
                    await asyncio.sleep(DELAY_ON_CONNECT_EXCEPTION_SECONDS)
            else:
                logger.warning(
                    f'Max retries exceeded ({RESULT_SENDER_MAX_RETRIES})'
                    f' for sending results of task {result.options.task_id}',
                )

            if not self.run_event.is_set():
                logger.info(
                    f'Current processing status: running={self.running_tasks} '
                    f'results={self.results_queue.qsize()}',
                )

    async def stopping(self):
        result_timeout = self.config.get_timeout('background') + RESULT_SENDER_WAIT_GRACE_SECONDS
        try:
            await asyncio.wait_for(
                self.results_task,
                timeout=result_timeout,
            )
        except asyncio.TimeoutError:
            logger.error(
                f'Cannot send all results timeout of {result_timeout} exceeded, cancel task',
            )
            self.results_task.cancel()
            try:
                await self.results_task
            except asyncio.CancelledError:
                logger.info('Result sender task has been cancelled')

    async def start(self):
        """
        Start the runner.
        """
        self.results_task = asyncio.create_task(self.result_sender())
        await super().start()
