# 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 queue and worker pool summaries."""

from typing import ClassVar, override

from django.test import override_settings

from debusine.artifacts.models import TaskTypes
from debusine.db.models import Artifact, WorkRequest, Worker
from debusine.server.status import (
    QueueStatus,
    TaskQueueSummary,
    WorkerStatus,
    WorkerSummary,
)
from debusine.server.tasks.wait.models import ExternalDebsignData
from debusine.server.workflows.models import WorkRequestWorkflowData
from debusine.tasks.models import (
    MmDebstrapBootstrapOptions,
    MmDebstrapData,
    SystemBootstrapRepository,
    WorkerType,
)
from debusine.test.django import TestCase


# We want to be in control of the scheduling status of tasks, for these tests
@override_settings(DISABLE_AUTOMATIC_SCHEDULING=True)
class WorkerStatusTests(TestCase):
    """Tests for WorkerStatus."""

    source: ClassVar[Artifact]
    environment: ClassVar[Artifact]

    @override
    @classmethod
    def setUpTestData(cls) -> None:
        super().setUpTestData()
        cls.source = cls.playground.create_minimal_source_package_artifact()
        environment = cls.playground.create_debian_environment()
        assert environment.artifact is not None
        cls.environment = environment.artifact

    def create_worker_task(
        self,
        architecture: str | None = None,
        *,
        mark_pending: bool = False,
        mark_running: bool = False,
        worker: Worker | None = None,
    ) -> WorkRequest:
        """Create a worker task for the given architecture."""
        if architecture is None:
            return self.playground.create_worker_task(
                task_name="noop",
                mark_pending=mark_pending,
                mark_running=mark_running,
                worker=worker,
            )
        else:
            work_request = self.playground.create_sbuild_work_request(
                source=self.source,
                environment=self.environment,
                architecture=architecture,
            )
            self.playground.advance_work_request(
                work_request,
                mark_pending=mark_pending,
                mark_running=mark_running,
                worker=worker,
            )
            return work_request

    def assertTasks(
        self, status: QueueStatus, task_type: TaskTypes, state: str, value: int
    ) -> None:
        """Assert tasks in status for task_type have state == value."""
        type_tasks = status.tasks.get(task_type, TaskQueueSummary())
        self.assertEqual(getattr(type_tasks, state), value)

    def assertWorkers(
        self,
        status: QueueStatus,
        worker_type: WorkerType,
        state: str,
        value: int,
    ) -> None:
        """Assert workers in status for worker_type have state == value."""
        type_status = status.workers.get(worker_type, WorkerSummary())
        self.assertEqual(getattr(type_status, state), value)

    def assertWorkerTasks(
        self,
        status: QueueStatus,
        worker_type: WorkerType,
        state: str,
        value: int,
    ) -> None:
        """Assert worker_tasks in status for worker_type have state == value."""
        type_tasks = status.worker_tasks.get(worker_type, TaskQueueSummary())
        self.assertEqual(getattr(type_tasks, state), value)

    def assertExternalWorkersArch(
        self, status: QueueStatus, arch: str, state: str, value: int
    ) -> None:
        """Assert external_workers_arch for arch have state == value."""
        arch_status = status.external_workers_arch.get(arch, WorkerSummary())
        self.assertEqual(getattr(arch_status, state), value)

    def assertExternalWorkersNativeArch(
        self, status: QueueStatus, arch: str, state: str, value: int
    ) -> None:
        """Assert external_workers_native_arch for arch have state == value."""
        arch_status = status.external_workers_native_arch.get(
            arch, WorkerSummary()
        )
        self.assertEqual(getattr(arch_status, state), value)

    def assertWorkerTasksArch(
        self, status: QueueStatus, arch: str, state: str, value: int
    ) -> None:
        """Assert worker_tasks_arch in status for arch have state == value."""
        arch_tasks = status.worker_tasks_arch.get(arch, TaskQueueSummary())
        self.assertEqual(getattr(arch_tasks, state), value)

    def test_counts_tasks_by_type(self) -> None:
        """Test that a tasks counts tasks by type."""
        self.create_worker_task(mark_pending=True)
        self.playground.create_signing_task(task_name="noop", mark_running=True)

        status = WorkerStatus.get_status()
        self.assertTasks(status, TaskTypes.WORKER, "pending", 1)
        self.assertTasks(status, TaskTypes.WORKER, "running", 0)
        self.assertTasks(status, TaskTypes.SIGNING, "pending", 0)
        self.assertTasks(status, TaskTypes.SIGNING, "running", 1)
        self.assertTasks(status, TaskTypes.INTERNAL, "pending", 0)
        self.assertTasks(status, TaskTypes.INTERNAL, "running", 0)
        self.assertWorkerTasks(status, WorkerType.EXTERNAL, "pending", 1)
        self.assertWorkerTasks(status, WorkerType.EXTERNAL, "running", 0)
        self.assertWorkerTasks(status, WorkerType.SIGNING, "pending", 0)
        self.assertWorkerTasks(status, WorkerType.SIGNING, "running", 1)
        self.assertWorkerTasks(status, WorkerType.CELERY, "pending", 0)
        self.assertWorkerTasks(status, WorkerType.CELERY, "running", 0)

    def test_worker_counts_against_all_archs(self) -> None:
        """Test that a worker counts against all architectures it supports."""
        worker = self.playground.create_worker(
            worker_type=WorkerType.EXTERNAL,
            extra_dynamic_metadata={
                "system:native_architecture": "amd64",
                "system:architectures": ["amd64", "i386"],
            },
        )
        worker.mark_connected()

        status = WorkerStatus.get_status()
        self.assertExternalWorkersArch(status, "amd64", "connected", 1)
        self.assertExternalWorkersArch(status, "i386", "connected", 1)
        # and by native_arch:
        self.assertExternalWorkersNativeArch(status, "amd64", "connected", 1)
        self.assertExternalWorkersNativeArch(status, "i386", "connected", 0)
        # and by type:
        self.assertWorkers(status, WorkerType.EXTERNAL, "connected", 1)

    def test_work_request_filters_by_task_type(self) -> None:
        """Test that worker_tasks_arch only includes WORKER work requests."""

    def test_work_request_counts_build_architecture_method(self) -> None:
        """
        Test that a work request is counted against its build_architecture.

        For a WorkRequest with build_architecture not specified in task_data.
        Instead the Tasks's .build_architecture() method is called to determine
        the architecture.
        """
        task_data = MmDebstrapData(
            bootstrap_options=MmDebstrapBootstrapOptions(architecture="amd64"),
            bootstrap_repositories=[
                SystemBootstrapRepository(
                    mirror="https://deb.debian.org/debian", suite="stable"
                )
            ],
        )
        expected_architecture = task_data.bootstrap_options.architecture

        self.playground.create_worker_task(
            task_name="mmdebstrap", task_data=task_data, mark_pending=True
        )
        self.playground.create_worker_task(
            task_name="mmdebstrap", task_data=task_data, mark_running=True
        )

        status = WorkerStatus.get_status()
        self.assertWorkerTasksArch(status, expected_architecture, "pending", 1)
        self.assertWorkerTasksArch(status, expected_architecture, "running", 1)
        self.assertTasks(status, TaskTypes.WORKER, "pending", 1)
        self.assertTasks(status, TaskTypes.WORKER, "running", 1)
        self.assertWorkerTasks(status, WorkerType.EXTERNAL, "pending", 1)
        self.assertWorkerTasks(status, WorkerType.EXTERNAL, "running", 1)

    def test_work_request_non_architecture_constrained(self) -> None:
        """
        Test that a work request without an architecture counts as arch:None.

        Work request without "build_architecture" in task_data and
        build_architecture() returns None.
        """
        self.create_worker_task(mark_pending=True)
        self.create_worker_task(mark_running=True)

        status = WorkerStatus.get_status()

        actual = status.worker_tasks_arch
        expected = {None: TaskQueueSummary(pending=1, running=1)}
        self.assertEqual(actual, expected)

        # Counts by type too:
        self.assertTasks(status, TaskTypes.WORKER, "pending", 1)
        self.assertTasks(status, TaskTypes.WORKER, "running", 1)
        self.assertWorkerTasks(status, WorkerType.EXTERNAL, "pending", 1)
        self.assertWorkerTasks(status, WorkerType.EXTERNAL, "running", 1)

    def test_get_status_no_pending_work_requests(self) -> None:
        """
        get_status: no pending work requests.

        Empty dictionary because no pending work requests and no workers.
        """
        status = WorkerStatus.get_status()
        self.assertEqual(status.external_workers_arch, {})
        self.assertEqual(status.worker_tasks_arch, {})
        for task_type in TaskTypes:
            self.assertEqual(status.tasks[task_type], TaskQueueSummary())
        for worker_type in WorkerType:
            if worker_type == WorkerType.SCHEDULER:
                self.assertNotIn(worker_type, status.workers)
                self.assertNotIn(worker_type, status.worker_tasks)
            else:
                self.assertEqual(status.workers[worker_type], WorkerSummary())
                self.assertEqual(
                    status.worker_tasks[worker_type], TaskQueueSummary()
                )

    def test_disabled_workers_ignored(self) -> None:
        """Workers are not counted if they are disabled."""
        self.playground.create_worker(
            extra_dynamic_metadata={
                "system:native_architecture": "arch1",
                "system:architectures": ["arch1", "arch2"],
            }
        )

        worker = self.playground.create_worker(
            extra_dynamic_metadata={
                "system:native_architecture": "arch3",
                "system:architectures": ["arch3"],
            }
        )
        assert worker.token is not None
        worker.token.disable()

        status = WorkerStatus.get_status()
        self.assertNotIn("arch3", status.external_workers_arch)
        self.assertNotIn("arch3", status.external_workers_native_arch)
        self.assertWorkers(status, WorkerType.EXTERNAL, "registered", 1)

    def test_count_registered_workers(self) -> None:
        """Registered and connected workers are counted."""
        worker = self.playground.create_worker(
            extra_dynamic_metadata={
                "system:native_architecture": "arch1",
                "system:architectures": ["arch1", "arch2"],
            }
        )
        worker.mark_connected()

        worker = self.playground.create_worker(
            extra_dynamic_metadata={
                "system:native_architecture": "arch3",
                "system:architectures": ["arch3"],
            }
        )

        status = WorkerStatus.get_status()
        self.assertExternalWorkersArch(status, "arch1", "registered", 1)
        self.assertExternalWorkersArch(status, "arch2", "registered", 1)
        self.assertExternalWorkersArch(status, "arch3", "registered", 1)
        self.assertExternalWorkersArch(status, "arch1", "connected", 1)
        self.assertExternalWorkersArch(status, "arch2", "connected", 1)
        self.assertExternalWorkersArch(status, "arch3", "connected", 0)
        # and by native_arch:
        self.assertExternalWorkersNativeArch(status, "arch1", "registered", 1)
        self.assertExternalWorkersNativeArch(status, "arch2", "registered", 0)
        self.assertExternalWorkersNativeArch(status, "arch3", "registered", 1)
        self.assertExternalWorkersNativeArch(status, "arch1", "connected", 1)
        self.assertExternalWorkersNativeArch(status, "arch2", "connected", 0)
        self.assertExternalWorkersNativeArch(status, "arch3", "connected", 0)
        # and by type:
        self.assertWorkers(status, WorkerType.EXTERNAL, "registered", 2)
        self.assertWorkers(status, WorkerType.EXTERNAL, "connected", 1)

    def test_count_registered_workers_by_type(self) -> None:
        """Registered workers are counted by type."""
        self.playground.create_worker(
            worker_type=WorkerType.EXTERNAL
        ).mark_connected()
        self.playground.create_worker(worker_type=WorkerType.EXTERNAL)
        self.playground.create_worker(
            worker_type=WorkerType.SIGNING
        ).mark_connected()
        self.playground.create_worker(
            worker_type=WorkerType.CELERY
        ).mark_connected()

        status = WorkerStatus.get_status()
        self.assertWorkers(status, WorkerType.EXTERNAL, "registered", 2)
        self.assertWorkers(status, WorkerType.SIGNING, "registered", 1)
        self.assertWorkers(status, WorkerType.CELERY, "registered", 1)
        self.assertWorkers(status, WorkerType.EXTERNAL, "connected", 1)
        self.assertWorkers(status, WorkerType.SIGNING, "connected", 1)
        self.assertWorkers(status, WorkerType.CELERY, "connected", 1)

    def test_count_busy_workers_running(self) -> None:
        """Assigned RUNNING tasks count as a busy worker."""
        worker = self.playground.create_worker(
            extra_dynamic_metadata={
                "system:native_architecture": "arch1",
                "system:architectures": ["arch1", "arch2"],
            }
        )
        worker.mark_connected()

        self.create_worker_task("arch1", worker=worker, mark_running=True)

        status = WorkerStatus.get_status()
        self.assertExternalWorkersArch(status, "arch1", "busy", 1)
        # the arch2 slot is also busy:
        self.assertExternalWorkersArch(status, "arch2", "busy", 1)
        # native_arch is busy:
        self.assertExternalWorkersNativeArch(status, "arch1", "busy", 1)
        # and by type:
        self.assertWorkers(status, WorkerType.EXTERNAL, "busy", 1)

    def test_count_busy_workers_pending(self) -> None:
        """Assigned PENDING tasks count as a busy worker."""
        worker = self.playground.create_worker(
            extra_dynamic_metadata={
                "system:native_architecture": "arch1",
                "system:architectures": ["arch1", "arch2"],
            }
        )
        worker.mark_connected()

        self.create_worker_task("arch1", mark_pending=True, worker=worker)

        status = WorkerStatus.get_status()
        self.assertExternalWorkersArch(status, "arch1", "busy", 1)
        # the arch2 slot is also busy:
        self.assertExternalWorkersArch(status, "arch2", "busy", 1)
        # build_arch is busy:
        self.assertExternalWorkersNativeArch(status, "arch1", "busy", 1)
        # and by type:
        self.assertWorkers(status, WorkerType.EXTERNAL, "busy", 1)

    def test_count_idle_workers(self) -> None:
        """Connected workers without assigned tasks are idle."""
        worker = self.playground.create_worker(
            extra_dynamic_metadata={
                "system:native_architecture": "arch1",
                "system:architectures": ["arch1"],
            }
        )
        worker.mark_connected()
        worker = self.playground.create_worker(
            extra_dynamic_metadata={
                "system:native_architecture": "arch2",
                "system:architectures": ["arch2"],
            }
        )
        status = WorkerStatus.get_status()
        self.assertExternalWorkersArch(status, "arch1", "idle", 1)
        self.assertExternalWorkersArch(status, "arch2", "registered", 1)
        self.assertExternalWorkersArch(status, "arch2", "idle", 0)
        # build_arch should be identical:
        self.assertExternalWorkersNativeArch(status, "arch1", "idle", 1)
        self.assertExternalWorkersNativeArch(status, "arch2", "registered", 1)
        self.assertExternalWorkersNativeArch(status, "arch2", "idle", 0)
        # and by type:
        self.assertWorkers(status, WorkerType.EXTERNAL, "registered", 2)
        self.assertWorkers(status, WorkerType.EXTERNAL, "idle", 1)

    def test_external_workers_arch_populated_for_all_task_archs(self) -> None:
        """We populate external_workers_arch for all queued architectures."""
        self.create_worker_task("arch1", mark_pending=True)
        status = WorkerStatus.get_status()
        self.assertIn("arch1", status.external_workers_arch)

    def test_worker_tasks_populated_for_all_worker_archs(self) -> None:
        """We populate worker_tasks_arch for all known worker architectures."""
        self.playground.create_worker(
            extra_dynamic_metadata={"system:architectures": ["arch1"]}
        )
        status = WorkerStatus.get_status()
        self.assertIn("arch1", status.worker_tasks_arch)

    def test_missing_system_architecture(self) -> None:
        """Workers without system:architecture log a warning."""
        self.playground.create_worker().mark_connected()
        with self.assertLogs("debusine.server.status") as cm:
            WorkerStatus.get_status()
        self.assertEqual(
            cm.output,
            [
                "WARNING:debusine.server.status:Worker computer-lan missing "
                "system:native_architecture"
            ],
        )

    def test_full_example(self) -> None:
        """Test the full data output of a non-trivial example state."""
        # Create three work requests of type sbuild
        self.create_worker_task("amd64", mark_pending=True)
        self.create_worker_task("amd64", mark_pending=True)
        self.create_worker_task("arm64", mark_pending=True)

        # Create three connected workers
        worker_amd64_1 = self.playground.create_worker(
            extra_dynamic_metadata={
                "system:native_architecture": "amd64",
                "system:architectures": ["amd64", "i386"],
            }
        )
        worker_amd64_1.mark_connected()

        worker_arm64_1 = self.playground.create_worker(
            extra_dynamic_metadata={
                "system:native_architecture": "arm64",
                "system:architectures": ["arm64"],
            }
        )
        worker_arm64_1.mark_connected()

        worker_arm64_2 = self.playground.create_worker(
            extra_dynamic_metadata={
                "system:native_architecture": "arm64",
                "system:architectures": ["arm64"],
            }
        )
        worker_arm64_2.mark_connected()

        # Create non-connected worker (will not appear as available
        # worker in the count for arm64)
        self.playground.create_worker(
            extra_dynamic_metadata={
                "system:native_architecture": "arm64",
                "system:architectures": ["arm64"],
            }
        )
        # And a celery worker
        self.playground.create_worker(
            worker_type=WorkerType.CELERY,
        ).mark_connected()
        # And a scheduler worker, which should not appear in the status
        self.playground.create_worker(
            worker_type=WorkerType.SCHEDULER
        ).mark_connected()

        # Create a running task
        self.create_worker_task(
            "arm64", mark_running=True, worker=worker_arm64_1
        )

        # Create a running wait task
        unsigned = self.playground.create_upload_artifacts().upload
        self.playground.create_wait_task(
            task_name="externaldebsign",
            task_data=ExternalDebsignData(unsigned=unsigned.id),
            workflow_data=WorkRequestWorkflowData(needs_input=True),
            mark_running=True,
        )

        actual = WorkerStatus.get_status()

        expected = QueueStatus(
            tasks={
                TaskTypes.WORKER: TaskQueueSummary(pending=3, running=1),
                TaskTypes.SERVER: TaskQueueSummary(),
                TaskTypes.INTERNAL: TaskQueueSummary(),
                TaskTypes.WORKFLOW: TaskQueueSummary(),
                TaskTypes.SIGNING: TaskQueueSummary(),
                TaskTypes.WAIT: TaskQueueSummary(running=1),
            },
            workers={
                WorkerType.CELERY: WorkerSummary(
                    registered=1, connected=1, idle=1
                ),
                WorkerType.EXTERNAL: WorkerSummary(
                    registered=4, connected=3, idle=2, busy=1
                ),
                WorkerType.SIGNING: WorkerSummary(),
            },
            worker_tasks={
                WorkerType.CELERY: TaskQueueSummary(),
                WorkerType.EXTERNAL: TaskQueueSummary(pending=3, running=1),
                WorkerType.SIGNING: TaskQueueSummary(),
            },
            external_workers_arch={
                "amd64": WorkerSummary(
                    registered=1, connected=1, idle=1, busy=0
                ),
                "arm64": WorkerSummary(
                    registered=3, connected=2, idle=1, busy=1
                ),
                "i386": WorkerSummary(
                    registered=1, connected=1, idle=1, busy=0
                ),
            },
            external_workers_native_arch={
                "amd64": WorkerSummary(
                    registered=1, connected=1, idle=1, busy=0
                ),
                "arm64": WorkerSummary(
                    registered=3, connected=2, idle=1, busy=1
                ),
            },
            worker_tasks_arch={
                "amd64": TaskQueueSummary(pending=2, running=0),
                "arm64": TaskQueueSummary(pending=1, running=1),
                "i386": TaskQueueSummary(pending=0, running=0),
            },
        )

        self.assertEqual(actual, expected)
