# 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.

"""Tests for TaskRunner."""

import asyncio
import datetime as dt
import logging
import queue
import stat
import time as tm
from collections.abc import Awaitable, Collection
from pathlib import Path
from typing import Literal, override
from unittest import mock

import debusine.tasks.models as task_models
from debusine.artifacts.models import (
    RuntimeStatistics,
    TaskTypes,
    WorkRequestResults,
)
from debusine.client.models import WorkRequestResponse
from debusine.tasks import BaseExternalTask, BaseTaskWithExecutor
from debusine.tasks.executors import (
    ExecutorInterface,
    ExecutorStatistics,
    InstanceInterface,
)
from debusine.tasks.models import WorkerType
from debusine.tasks.tests.helpers import SampleBaseExternalTask
from debusine.test import TestCase
from debusine.test.test_utils import (
    create_work_request_response,
    preserve_task_registry,
)
from debusine.worker.config import ConfigHandler
from debusine.worker.system_information import cpu_time
from debusine.worker.task_runner import (
    CollectStatistics,
    TaskRunner,
    TaskRunnerError,
)
from debusine.worker.tests import server


class WorkerConfigTestCase(TestCase, server.DebusineAioHTTPTestCase):
    """Test case with a worker configuration."""

    @override
    async def setUpAsync(self) -> None:
        await super().setUpAsync()
        self.api_url = str(self.server.make_url('')) + "/api"

        self.temp_config_directory = self.create_temp_config_directory(
            {
                'General': {'api-url': self.api_url},
                'Environment': {'ENABLE': 'yes'},
            },
        )

        self.config = ConfigHandler(
            directories=[self.temp_config_directory.as_posix()]
        )
        self.config['General']['log-file'] = str(
            self.temp_config_directory / 'worker.log'
        )


class CollectStatisticsTests(TestCase, server.DebusineAioHTTPTestCase):
    """Tests for :py:class:`CollectStatistics`."""

    @override
    async def setUpAsync(self) -> None:
        await super().setUpAsync()
        self.statistics = RuntimeStatistics()
        self.task = mock.MagicMock(spec=BaseTaskWithExecutor)
        self.task.executor = None

    def set_executor(self, statistics: ExecutorStatistics) -> None:
        self.task.executor = mock.MagicMock(spec=ExecutorInterface)
        self.task.executor.get_statistics.return_value = statistics
        self.task.executor_instance = mock.MagicMock(spec=InstanceInterface)

    def test_init(self) -> None:
        collect = CollectStatistics(self.task, self.statistics)
        self.assertIs(collect.task, self.task)
        self.assertIs(collect.statistics, self.statistics)
        self.assertLessEqual(collect.start_time, tm.monotonic())
        self.assertLessEqual(collect.start_cpu_time, cpu_time())
        self.assertFalse(collect.done.is_set())
        self.assertIsNone(collect.collect_task)

    def test_gather_no_executor(self) -> None:
        collect = CollectStatistics(self.task, self.statistics)
        collect.gather()
        self.assertEqual(self.statistics, RuntimeStatistics())

    def test_gather(self) -> None:
        es = ExecutorStatistics(
            disk_space=1,
            memory=2,
            available_disk_space=4,
            available_memory=8,
            cpu_count=16,
        )
        self.set_executor(es)
        collect = CollectStatistics(self.task, self.statistics)
        with self.spy(
            "debusine.tasks.executors.base.ExecutorStatistics.add_to"
        ) as add_to:
            collect.gather()
        self.assertEqual(
            self.statistics,
            RuntimeStatistics(
                disk_space=1,
                memory=2,
                available_disk_space=4,
                available_memory=8,
                cpu_count=16,
            ),
        )
        add_to.assert_called_once_with(es, self.statistics)

    async def test_collect(self) -> None:
        collect = CollectStatistics(self.task, self.statistics)
        self.assertIsNone(self.statistics.duration)
        self.assertIsNone(self.statistics.cpu_time)

        gather = self.enterContext(mock.patch.object(collect, "gather"))

        count = 0

        async def mock_wait(
            _aws: Collection[Awaitable[None]],
            timeout: int,  # noqa: ASYNC109, ARG001
        ) -> None:
            nonlocal count
            count += 1
            if count == 3:
                collect.done.set()

        self.enterContext(mock.patch("asyncio.wait", side_effect=mock_wait))

        await collect.collect()

        self.assertEqual(gather.call_count, 3)

        self.assertIsNotNone(self.statistics.duration)
        self.assertIsNotNone(self.statistics.cpu_time)

    async def test_context_manager(self) -> None:
        collect = CollectStatistics(self.task, self.statistics)
        gather = self.enterContext(mock.patch.object(collect, "gather"))

        self.assertIsNone(collect.collect_task)
        self.assertFalse(collect.done.is_set())

        async with collect:
            self.assertIsNotNone(collect.collect_task)

        self.assertGreaterEqual(gather.call_count, 1)

        self.assertTrue(collect.done.is_set())
        self.assertIsNone(collect.collect_task)


class TaskRunnerTests(WorkerConfigTestCase):
    """Tests for :py:class:`TaskRunner`."""

    runner: TaskRunner

    @override
    async def setUpAsync(self) -> None:
        await super().setUpAsync()
        self.workdir = self.config.working_directory
        self.workdir.mkdir(parents=True, exist_ok=True)
        self.worker_type: Literal[WorkerType.EXTERNAL, WorkerType.SIGNING] = (
            WorkerType.EXTERNAL
        )
        self.work_request = WorkRequestResponse(
            id=42,
            url=self.api_url + "/work_requests/42",
            created_at=dt.datetime(2026, 1, 1, 0, tzinfo=dt.UTC),
            started_at=dt.datetime(2026, 1, 1, 1, tzinfo=dt.UTC),
            status="running",
            result="",
            worker=1,
            task_type="Worker",
            task_name="noop",
            task_data={},
            token="work-request-token",
            dynamic_task_data={},
            priority_base=0,
            priority_adjustment=0,
            artifacts=[],
            scope="debusine",
            workspace="Playground",
        )

    @override
    async def tearDownAsync(self) -> None:
        await self.runner.shutdown()
        await super().tearDownAsync()

    def setup_runner(self) -> None:
        """Set up a runner for this work request."""
        self.runner = TaskRunner(
            self.config, self.worker_type, self.work_request, self.workdir
        )

    def test_init(self) -> None:
        self.setup_runner()
        self.assertIs(self.runner.config, self.config)
        self.assertIs(self.runner.worker_type, self.worker_type)
        self.assertIs(self.runner.work_request, self.work_request)
        self.assertTrue(self.runner.workdir.is_relative_to(self.workdir))
        st = self.runner.workdir.stat()
        self.assertEqual(stat.S_IMODE(st.st_mode), 0o755)
        self.assertEqual(self.runner.client.base_api_url, self.config.api_url)
        self.assertEqual(
            self.runner.client.token, self.runner.work_request.token
        )
        self.assertEqual(self.runner.client.timeout, self.config.http_timeout)
        self.assertEqual(self.runner.client.scope, self.work_request.scope)
        self.assertIsNone(self.runner.task)

    async def test_setup_task(self) -> None:
        set_working_directory = self.enterContext(
            mock.patch.object(BaseExternalTask, "set_working_directory")
        )
        set_worker_environment_variables = self.enterContext(
            mock.patch.object(
                BaseExternalTask, "set_worker_environment_variables"
            )
        )
        configure_server_access = self.enterContext(
            mock.patch.object(BaseExternalTask, "configure_server_access")
        )

        self.setup_runner()
        await self.runner.setup_task()
        task = self.runner.task
        assert task is not None

        set_working_directory.assert_called_once_with(self.runner.workdir)
        set_worker_environment_variables.assert_called_once_with(
            {"ENABLE": "yes"}
        )
        configure_server_access.assert_called_once_with(self.runner.client)
        self.assertEqual(task.work_request_id, self.work_request.id)
        self.assertEqual(task.workspace_name, self.work_request.workspace)

    async def test_setup_task_signing_task_on_external_worker(self) -> None:
        self.enterContext(preserve_task_registry())

        class SigningTestData(task_models.BaseTaskData):
            pass

        class SigningTest(
            SampleBaseExternalTask[
                SigningTestData, task_models.BaseDynamicTaskData
            ],
        ):
            TASK_VERSION = 1
            TASK_TYPE = TaskTypes.SIGNING

            @override
            def run(self, execute_directory: Path) -> WorkRequestResults:  # noqa: U100
                raise NotImplementedError()

        self.work_request = create_work_request_response(
            id=52,
            task_type="Signing",
            task_name="signingtest",
            created_at=dt.datetime.fromisoformat("2022-01-05T11:14:22.242178Z"),
            started_at=None,
            completed_at=None,
            duration=None,
            worker=58,
            task_data={},
            token="work-request-token",
            dynamic_task_data={},
            priority_base=0,
            priority_adjustment=0,
            status="running",
            result="",
            artifacts=[],
            workspace="TestWorkspace",
        )
        self.setup_runner()
        with self.assertRaises(TaskRunnerError) as exc:
            await self.runner.setup_task()

        self.assertEqual(exc.exception.code, "wrong-task-type")
        self.assertEqual(
            exc.exception.message, "Task: %s is of type %s, not %s"
        )
        self.assertEqual(
            exc.exception.args,
            (
                self.work_request.task_name,
                self.work_request.task_type,
                "Worker",
            ),
        )
        self.assertIsNone(self.runner.task)

    async def test_setup_task_worker_task_on_signing_worker(self) -> None:
        self.worker_type = WorkerType.SIGNING
        self.setup_runner()
        with self.assertRaises(TaskRunnerError) as exc:
            await self.runner.setup_task()

        self.assertEqual(exc.exception.code, "wrong-task-type")
        self.assertEqual(
            exc.exception.message, "Task: %s is of type %s, not %s"
        )
        self.assertEqual(
            exc.exception.args,
            (
                self.work_request.task_name,
                self.work_request.task_type,
                "Signing",
            ),
        )
        self.assertIsNone(self.runner.task)

    async def test_setup_task_signing_task_on_signing_worker(self) -> None:
        self.enterContext(preserve_task_registry())

        class SigningTestData(task_models.BaseTaskData):
            pass

        class SigningTest(
            SampleBaseExternalTask[
                SigningTestData, task_models.BaseDynamicTaskData
            ],
        ):
            TASK_VERSION = 1
            TASK_TYPE = TaskTypes.SIGNING

            @override
            def run(self, execute_directory: Path) -> WorkRequestResults:  # noqa: U100
                raise NotImplementedError()

        self.worker_type = WorkerType.SIGNING
        self.work_request = create_work_request_response(
            id=52,
            task_type="Signing",
            task_name="signingtest",
            created_at=dt.datetime.fromisoformat("2022-01-05T11:14:22.242178Z"),
            started_at=None,
            completed_at=None,
            duration=None,
            worker=58,
            task_data={},
            token="work-request-token",
            dynamic_task_data={},
            priority_base=0,
            priority_adjustment=0,
            status="running",
            result="",
            artifacts=[],
            workspace="TestWorkspace",
        )
        self.setup_runner()
        await self.runner.setup_task()
        self.assertIsNotNone(self.runner.task)

    async def test_setup_task_server_task_on_external_worker(self) -> None:
        self.enterContext(preserve_task_registry())

        class ServerTestData(task_models.BaseTaskData):
            pass

        class ServerTest(
            SampleBaseExternalTask[
                ServerTestData, task_models.BaseDynamicTaskData
            ],
        ):
            TASK_VERSION = 1
            TASK_TYPE = TaskTypes.SERVER

            @override
            def _execute(self) -> WorkRequestResults:
                raise NotImplementedError()

            @override
            def _upload_work_request_debug_logs(self) -> None:
                raise NotImplementedError()

        self.work_request = create_work_request_response(
            id=52,
            task_type="Server",
            task_name="servertest",
            created_at=dt.datetime.fromisoformat("2022-01-05T11:14:22.242178Z"),
            started_at=None,
            completed_at=None,
            duration=None,
            worker=58,
            task_data={},
            token="work-request-token",
            dynamic_task_data={},
            priority_base=0,
            priority_adjustment=0,
            status="running",
            result="",
            artifacts=[],
            workspace="TestWorkspace",
        )

        self.setup_runner()
        with self.assertRaises(TaskRunnerError) as exc:
            await self.runner.setup_task()

        self.assertEqual(exc.exception.code, "wrong-task-type")
        self.assertEqual(
            exc.exception.message, "Task: %s is of type %s, not %s"
        )
        self.assertEqual(
            exc.exception.args,
            (
                self.work_request.task_name,
                self.work_request.task_type,
                "Worker",
            ),
        )
        self.assertIsNone(self.runner.task)

    async def test_setup_task_unknown_task_name(self) -> None:
        self.work_request = create_work_request_response(
            id=52,
            task_type="Worker",
            task_name="nonexistent",
            created_at=dt.datetime.fromisoformat("2022-01-05T11:14:22.242178Z"),
            started_at=None,
            completed_at=None,
            duration=None,
            worker=58,
            task_data={},
            token="work-request-token",
            dynamic_task_data={},
            priority_base=0,
            priority_adjustment=0,
            status="running",
            result="",
            artifacts=[],
            workspace="TestWorkspace",
        )
        self.setup_runner()
        with self.assertRaises(TaskRunnerError) as exc:
            await self.runner.setup_task()

        self.assertEqual(exc.exception.code, "setup-failed")
        self.assertEqual(exc.exception.message, "Task: %s Error setup: %s")
        self.assertEqual(exc.exception.args[0], self.work_request.task_name)
        self.assertEqual(
            str(exc.exception.args[1]),
            "'nonexistent' is not a registered Worker task_name",
        )
        self.assertIsNone(self.runner.task)

    async def test_setup_task_invalid_task_data(self) -> None:
        self.work_request.task_data = {"invalid": True}
        self.setup_runner()
        with self.assertRaises(TaskRunnerError) as exc:
            await self.runner.setup_task()

        self.assertEqual(exc.exception.code, "configure-failed")
        self.assertEqual(exc.exception.message, "Task: %s Error configure: %s")
        self.assertEqual(exc.exception.args[0], self.work_request.task_name)
        self.assertRegex(
            str(exc.exception.args[1]), r"Extra inputs are not permitted"
        )
        self.assertIsNone(self.runner.task)

    async def test_execute(self) -> None:
        gather = self.enterContext(
            mock.patch.object(CollectStatistics, "gather")
        )
        execute = self.enterContext(
            mock.patch(
                "debusine.tasks.Noop.execute_logging_exceptions",
                return_value=WorkRequestResults.SUCCESS,
            )
        )
        self.setup_runner()
        await self.runner.setup_task()
        self.assertEqual(
            await self.runner.execute(), WorkRequestResults.SUCCESS
        )

        gather.assert_called()
        execute.assert_called()

    async def test_execute_raises(self) -> None:
        gather = self.enterContext(
            mock.patch.object(CollectStatistics, "gather")
        )
        self.enterContext(
            mock.patch(
                "debusine.tasks.Noop.execute_logging_exceptions",
                side_effect=RuntimeError("expected error"),
            )
        )
        self.setup_runner()
        await self.runner.setup_task()
        with self.assertRaises(TaskRunnerError) as exc:
            await self.runner.execute()

        gather.assert_called()

        self.assertEqual(exc.exception.code, "execute-failed")
        self.assertEqual(exc.exception.message, "Task: %s Error execute: %s")
        self.assertEqual(exc.exception.args[0], self.work_request.task_name)
        self.assertEqual(str(exc.exception.args[1]), "expected error")

    async def test_cancel_execution_aborts_task(self) -> None:
        self.enterContext(mock.patch.object(CollectStatistics, "gather"))

        loop = asyncio.get_running_loop()
        executing = asyncio.Event()
        test_done: queue.Queue[None] = queue.Queue()

        def mock_task_execute() -> WorkRequestResults:
            nonlocal executing
            loop.call_soon_threadsafe(executing.set)
            test_done.get()
            return WorkRequestResults.SUCCESS

        self.enterContext(
            mock.patch(
                "debusine.tasks.noop.Noop.execute_logging_exceptions",
                side_effect=mock_task_execute,
            )
        )
        self.setup_runner()
        await self.runner.setup_task()

        # Wait until the runner starts executing the task
        execution = asyncio.create_task(self.runner.execute())
        await executing.wait()
        assert self.runner.task is not None
        self.assertFalse(self.runner.task.aborted)

        # Cancel the task execution
        execution.cancel()
        await asyncio.wait([execution])

        self.assertTrue(self.runner.task.aborted)
        test_done.put_nowait(None)

    async def test_context_manager(self) -> None:
        gather = self.enterContext(
            mock.patch.object(CollectStatistics, "gather")
        )
        execute = self.enterContext(
            mock.patch(
                "debusine.tasks.Noop.execute_logging_exceptions",
                return_value=WorkRequestResults.SUCCESS,
            )
        )
        self.setup_runner()

        self.assertIsNone(self.runner.task)
        async with self.runner:
            await self.runner.setup_task()
            self.assertIsNotNone(self.runner.task)
            self.assertTrue(self.runner.workdir.is_dir())
            self.assertEqual(
                await self.runner.execute(), WorkRequestResults.SUCCESS
            )
        self.assertFalse(self.runner.workdir.is_dir())

        gather.assert_called()
        execute.assert_called()

    async def test_shutdown_cleanup_error(self) -> None:
        self.setup_runner()

        with self.assertLogs(
            logger="debusine.worker.task_runner", level=logging.WARNING
        ) as log:
            async with self.runner:
                with mock.patch.object(
                    self.runner.workdir_tempdir,
                    "cleanup",
                    side_effect=RuntimeError("expected error"),
                ):
                    await self.runner.shutdown()

        [warning] = log.output
        self.assertRegex(
            warning,
            r"WARNING:debusine.worker.task_runner:"
            "Failed to remove '/.+/WR-42-.+':"
            r" expected error",
        )
