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

"""Unit tests for base server-side task classes."""

from pathlib import Path

from django.test import TestCase

from debusine.artifacts.local_artifact import WorkRequestDebugLogs
from debusine.db.models import Artifact, ArtifactRelation, WorkRequest
from debusine.server.tasks.base import BaseServerTask
from debusine.tasks.models import BaseDynamicTaskData, BaseTaskData, WorkerType
from debusine.tasks.tests.helper_mixin import TestTaskMixin
from debusine.tasks.tests.test_task import TestBaseTask2Data
from debusine.test import TestHelpersMixin
from debusine.worker.system_information import host_architecture


class TestBaseServerTask(
    TestTaskMixin, BaseServerTask[BaseTaskData, BaseDynamicTaskData]
):
    """Sample class to test BaseServerTask class."""

    def _execute(self) -> bool:
        """Unused abstract method from BaseTask."""
        raise NotImplementedError()


class TestBaseServerTask2(
    TestTaskMixin, BaseServerTask[TestBaseTask2Data, BaseDynamicTaskData]
):
    """Test BaseServerTask class with jsonschema validation."""

    TASK_VERSION = 1

    def _execute(self) -> bool:
        """Unused abstract method from BaseTask."""
        raise NotImplementedError()


class BaseServerTaskTests(TestHelpersMixin, TestCase):
    """Unit tests for :class:`BaseServerTask`."""

    def setUp(self) -> None:
        """Create the shared attributes."""
        self.task = TestBaseServerTask({})
        self.task2 = TestBaseServerTask2({"foo": "bar"})
        self.worker_metadata = {"system:worker_type": WorkerType.CELERY}

    def test_can_run_on_no_version(self):
        """Ensure can_run_on returns True if no version is specified."""
        self.assertIsNone(self.task.TASK_VERSION)
        metadata = {**self.worker_metadata, **self.task.analyze_worker()}
        self.assertEqual(self.task.can_run_on(metadata), True)

    def test_can_run_on_with_different_versions(self):
        """Ensure can_run_on returns False if versions differ."""
        self.assertIsNone(self.task.TASK_VERSION)
        metadata = {**self.worker_metadata, **self.task.analyze_worker()}
        metadata["server:testbaseservertask:version"] = 1
        self.assertEqual(self.task.can_run_on(metadata), False)

    def test_set_work_request(self) -> None:
        """set_work_request sets appropriate attributes."""
        work_request = self.playground.create_work_request()

        self.task.set_work_request(work_request)

        self.assertEqual(self.task.work_request, work_request)
        self.assertEqual(self.task.workspace, work_request.workspace)
        self.assertEqual(self.task.work_request_id, work_request.id)
        self.assertEqual(self.task.workspace_name, work_request.workspace.name)
        self.assertEqual(
            self.task.worker_host_architecture, host_architecture()
        )

    def setup_upload_work_request_debug_logs(
        self, source_artifacts: list[Artifact] | None = None
    ):
        """Setup for upload_work_request_debug_logs tests."""  # noqa: D401
        # Add a file to be uploaded
        with self.task.open_debug_log_file("test.log") as file:
            file.write("log")
        assert self.task._debug_log_files_directory is not None

        self.task.set_work_request(self.playground.create_work_request())
        self.task._source_artifacts_ids = [
            artifact.id for artifact in source_artifacts or []
        ]

        return [
            (path.name, path.read_bytes())
            for path in Path(
                self.task._debug_log_files_directory.name
            ).iterdir()
        ]

    def assert_uploaded_work_request_debug_logs_artifact(
        self,
        work_request: WorkRequest,
        expected_files: list[tuple[str, bytes]],
    ) -> Artifact:
        """Assert that an Artifact was created."""
        artifact = Artifact.objects.filter(
            category=WorkRequestDebugLogs._category
        ).last()
        assert artifact is not None
        assert artifact.created_by_work_request is not None

        self.assertEqual(artifact.category, WorkRequestDebugLogs._category)
        self.assertEqual(artifact.workspace, work_request.workspace)
        self.assertEqual(artifact.data, {})
        self.assertEqual(artifact.created_by_work_request, work_request)
        file_store = (
            work_request.workspace.default_file_store.get_backend_object()
        )
        for file_in_artifact, (name, contents) in zip(
            artifact.fileinartifact_set.order_by("id"), expected_files
        ):
            self.assertEqual(file_in_artifact.path, name)
            assert file_in_artifact.file is not None
            with file_store.get_stream(file_in_artifact.file) as file:
                self.assertEqual(file.read(), contents)

        return artifact

    def test_upload_work_request_debug_logs_with_relation(self):
        """
        Artifact is created and uploaded. Relation is created.

        The relation is from the debug logs artifact to the source_artifact_id.
        """
        source_artifact = self.create_artifact()[0]
        expected_files = self.setup_upload_work_request_debug_logs(
            [source_artifact]
        )

        self.task._upload_work_request_debug_logs()

        artifact = self.assert_uploaded_work_request_debug_logs_artifact(
            self.task.work_request, expected_files
        )

        self.assertEqual(artifact.relations.count(), 1)
        relation = artifact.relations.first()
        assert relation is not None
        self.assertEqual(relation.artifact, artifact)
        self.assertEqual(relation.target, source_artifact)
        self.assertEqual(relation.type, ArtifactRelation.Relations.RELATES_TO)

    def test_upload_work_request_debug_logs(self):
        """Artifact is created and uploaded."""
        expected_files = self.setup_upload_work_request_debug_logs()

        self.task._upload_work_request_debug_logs()

        self.assert_uploaded_work_request_debug_logs_artifact(
            self.task.work_request, expected_files
        )

    def test_upload_work_request_no_log_files(self):
        """No log files: no artifact created."""
        self.task._upload_work_request_debug_logs()

        self.assertFalse(
            Artifact.objects.filter(
                category=WorkRequestDebugLogs._category
            ).exists()
        )
