# Copyright © The Debusine Developers
# See the AUTHORS file at the top-level directory of this distribution
#
# This file is part of Debusine. It is subject to the license terms
# in the LICENSE file found in the top-level directory of this
# distribution. No part of Debusine, including this file, may be copied,
# modified, propagated, or distributed except according to the terms
# contained in the LICENSE file.

"""Supervise running tasks for the worker."""

import asyncio
import logging
import tempfile
import time as tm
import types
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Any, Literal, Self, assert_never

from debusine.artifacts.models import (
    RuntimeStatistics,
    TaskTypes,
    WorkRequestResults,
)
from debusine.client.debusine import Debusine
from debusine.client.models import WorkRequestResponse
from debusine.tasks import BaseExternalTask, TaskConfigError
from debusine.tasks.models import WorkerType
from debusine.worker.config import ConfigHandler
from debusine.worker.system_information import cpu_time

logger = logging.getLogger(__name__)


class TaskRunnerError(Exception):
    """Report an error while running a task."""

    def __init__(self, code: str, message: str, *args: Any) -> None:
        """
        Initialize a TaskRunnerError.

        :param code: the error code
        :param message: the error message, suitable for logging, with
          percent-expansion placeholders
        :param args: arguments to use to expand placeholders in ``message``.
        """
        self.code = code
        self.message = message
        self.args = args


class CollectStatistics:
    """Handle collecting statistics for a task execution."""

    def __init__(
        self,
        task: BaseExternalTask[Any, Any],
        statistics: RuntimeStatistics,
    ) -> None:
        """Initialize a CollectStatistics."""
        self.task = task

        # Statistics collected from the worker and task executors.
        self.statistics = statistics

        # The monotonic clock time when the task started executing.
        self.start_time: float = tm.monotonic()

        # The total CPU time used so far when the task started executing.
        # (We assume that the worker only runs one task at a time.)
        self.start_cpu_time: float = cpu_time()

        # Set this event when the task is done
        self.done = asyncio.Event()

        # Collection task
        self.collect_task: asyncio.Task[None] | None = None

    def gather(self) -> None:
        """Perform one round of statistics collection from the executor."""
        # We can only start collecting statistics once an executor is
        # instantiated.
        if self.task.executor is None:
            return

        # TODO: this stops the main loop. The executor API probably needs to
        #       have this made async
        statistics = self.task.executor.get_statistics(
            self.task.executor_instance
        )
        statistics.add_to(self.statistics)

    async def __aenter__(self) -> Self:
        """Start statistics collection."""
        self.gather()
        assert self.collect_task is None
        self.collect_task = asyncio.create_task(self.collect())
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: types.TracebackType | None,
    ) -> None:
        """Stop statistics collection."""
        self.done.set()
        assert self.collect_task is not None
        await self.collect_task
        self.collect_task = None

    async def collect(self) -> None:
        """Collect statistics until self.done is set."""
        try:
            done_task = asyncio.create_task(self.done.wait())
            while not self.done.is_set():
                self.gather()

                # Arbitrary polling interval.  We want to poll often enough to
                # get useful information about relatively brief spikes in disk
                # or RAM use, but not so often as to spend excessive amounts of
                # time gathering statistics.
                await asyncio.wait([done_task], timeout=10)
        finally:
            self.statistics.duration = int(tm.monotonic() - self.start_time)
            self.statistics.cpu_time = int(cpu_time() - self.start_cpu_time)


class TaskRunner:
    """Supervises running a task."""

    def __init__(
        self,
        config: ConfigHandler,
        worker_type: Literal[WorkerType.EXTERNAL, WorkerType.SIGNING],
        work_request: WorkRequestResponse,
        worker_workdir: Path,
    ) -> None:
        """
        Initialize a TaskRunner.

        :param config: worker configuration
        :param worker_type: worker type
        :param work_request: work request to run
        :param worker_workdir: ``Path`` to the directory where our working
          directory should be created
        """
        self.config = config
        self.worker_type = worker_type
        self.work_request = work_request
        self.workdir_tempdir = tempfile.TemporaryDirectory(
            prefix=f"WR-{self.work_request.id}-", dir=worker_workdir
        )
        self.workdir = Path(self.workdir_tempdir.name)
        self.workdir.chmod(0o755)
        self.client = Debusine(
            self.config.api_url,
            work_request.token,
            timeout=self.config.http_timeout,
            logger=logger,
            scope=work_request.scope,
        )
        # Running task
        self.task: BaseExternalTask[Any, Any] | None = None

        # Statistics collected from the worker and task executors.
        self.statistics = RuntimeStatistics()

    async def setup_task(self) -> None:
        """Instantiate the task to run."""
        task_type = TaskTypes(self.work_request.task_type)
        match self.worker_type:
            case WorkerType.EXTERNAL:
                expected_task_type = TaskTypes.WORKER
            case WorkerType.SIGNING:
                expected_task_type = TaskTypes.SIGNING
            case _ as unreachable:
                assert_never(unreachable)

        if task_type != expected_task_type:
            raise TaskRunnerError(
                "wrong-task-type",
                "Task: %s is of type %s, not %s",
                self.work_request.task_name,
                task_type,
                expected_task_type,
            )

        try:
            task_class = BaseExternalTask.class_from_name(
                TaskTypes(self.work_request.task_type),
                self.work_request.task_name,
            )
        except ValueError as exc:
            raise TaskRunnerError(
                "setup-failed",
                "Task: %s Error setup: %s",
                self.work_request.task_name,
                exc,
            )

        try:
            self.task = task_class(
                self.work_request.task_data, self.work_request.dynamic_task_data
            )
            # Inject the worker's configured HTTP timeout into the task instance
            self.task.http_timeout = self.config.http_timeout
        except TaskConfigError as exc:
            raise TaskRunnerError(
                "configure-failed",
                "Task: %s Error configure: %s",
                self.work_request.task_name,
                exc,
            )

        assert self.task.TASK_TYPE == self.work_request.task_type
        assert isinstance(self.task, BaseExternalTask)
        self.task.set_working_directory(self.workdir)
        self.task.set_worker_environment_variables(self.config.environment)

        self.task.configure_server_access(self.client)
        self.task.work_request_id = self.work_request.id
        self.task.workspace_name = self.work_request.workspace

    async def shutdown(self) -> None:
        """Shut down after execution."""
        try:
            self.workdir_tempdir.cleanup()
        except Exception as exc:
            logger.warning("Failed to remove '%s': %s", self.workdir, str(exc))

    async def __aenter__(self) -> Self:
        """Set up running the task."""
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc_val: BaseException | None,
        exc_tb: types.TracebackType | None,
    ) -> None:
        """Cleanup after running the task."""
        await self.shutdown()

    async def execute(self) -> WorkRequestResults:
        """
        Run the task.

        :returns: a completed Future with the task results or an exception
        :raises TaskRunError: on errors
        """
        assert self.task is not None
        # ThreadPool for BaseExternalTask.execute (e.g. Sbuild.exec).
        # We only allow a single thread.  We only run one task at a time
        # anyway, and signing workers need to run Django code in tasks which is
        # safest if done from only one thread.
        #
        # TODO: it is impossible to interrupt a running task using
        # concurrent.futures executors. Given that we're mostly executing
        # external commands, we can migrate to something based on
        # asyncio.Subprocess
        task_executor = ThreadPoolExecutor(max_workers=1)

        try:
            async with CollectStatistics(self.task, self.statistics):
                task_exec_future = task_executor.submit(
                    self.task.execute_logging_exceptions
                )
                try:
                    return await asyncio.wrap_future(task_exec_future)
                except Exception as exc:
                    raise TaskRunnerError(
                        "execute-failed",
                        "Task: %s Error execute: %s",
                        self.task.name,
                        exc,
                    )
        except asyncio.CancelledError:
            self.task.abort()
            raise
        finally:
            task_executor.shutdown(wait=False, cancel_futures=True)
