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

"""Common test-helper code involving Tasks."""

import abc
import datetime as dt
import hashlib
import shlex
import tarfile
from io import BytesIO
from pathlib import Path
from shutil import rmtree
from tempfile import mkdtemp
from typing import Any, ClassVar, cast, override
from unittest import mock

import debusine.tasks.tags as ttags
import debusine.worker.tags as wtags
from debusine.artifacts.models import ArtifactCategory, WorkRequestResults
from debusine.client.debusine import Debusine
from debusine.client.models import (
    ArtifactResponse,
    FileResponse,
    FilesResponseType,
)
from debusine.tasks import BaseExternalTask, BaseTaskWithExecutor, inputs
from debusine.tasks.executors import ExecutorInterface, UnshareInstance
from debusine.tasks.inputs.mock_resolver import MockFieldResolver
from debusine.tasks.models import (
    BaseDynamicTaskData,
    BaseDynamicTaskDataWithExecutor,
    BaseTaskData,
    BaseTaskDataWithExecutor,
)
from debusine.test.base import TestCase
from debusine.test.test_utils import create_artifact_response
from debusine.utils import extract_generic_type_arguments


class CurrentTaskAccessor[ValueType]:
    """Define a class member that depends on the current task being set."""

    name: str

    def __set_name__(self, owner: type["TaskTestCase[Any]"], name: str) -> None:
        """Set the field name at class construction."""
        self.name = name

    def __get__(
        self,
        obj: "TaskTestCase[Any]",
        _: "type[TaskTestCase[Any]] | None" = None,
    ) -> ValueType:
        """Initialize the current task the first time the field is accessed."""
        obj.task_set_current(obj.task_create())
        return cast(ValueType, obj.__dict__[self.name])


class TaskTestCase[TaskClass](TestCase, abc.ABC):
    """
    Base test case for all task types.

    The task being tested is instantiated the first time it is accessed, so you
    can edit self.task_data before access to customize the task data for the
    current test.
    """

    #: Common task data for all tests
    SAMPLE_TASK_DATA: ClassVar[dict[str, Any]] = {}
    #: Test-specific task data. It is initialized with SAMPLE_TASK_DATA and can
    #: be edited freely
    task_data: dict[str, Any]

    #: Class for the task being tested
    task_class: type[TaskClass]

    #: Tasks instantiated by the current test
    tasks: list[TaskClass]

    task = CurrentTaskAccessor[TaskClass]()

    @override
    def __init_subclass__(cls, **kwargs: Any) -> None:
        """Compute cls.task_class."""
        super().__init_subclass__(**kwargs)

        # The task class used for this test suite, computed by introspecting
        # the type argument used to specialize this generic class.
        [cls.task_class] = extract_generic_type_arguments(cls, TaskTestCase)

    @override
    def setUp(self) -> None:
        super().setUp()
        self.task_data = self.SAMPLE_TASK_DATA.copy()
        self.tasks = []

    @override
    def tearDown(self) -> None:
        """Delete debug log files directory if it exists."""
        for task in self.tasks:
            self.task_teardown(task)
        super().tearDown()

    @abc.abstractmethod
    def task_create(self, **kwargs: Any) -> TaskClass:
        """
        Instantiate a task using the current task data.

        :param kwargs: if any keyword arguments are provided, they replace the
          corresponding values in ``self.task_data``.
        """

    def task_manage(self, task: TaskClass) -> None:
        """Register a task for task_setup and task_teardown."""
        self.task_setup(task)
        self.tasks.append(task)

    def task_setup(self, task: TaskClass) -> None:
        """Instantiate the task to test."""

    def task_teardown(self, task: TaskClass) -> None:
        """Perform cleanup at the end of the lifetime of the current task."""

    def task_reset_current(self) -> None:
        """Remove the current task, if already instantiated."""
        self.__dict__.pop("task", None)

    def task_set_current(self, task: TaskClass) -> None:
        """Set the given task as self.task."""
        if task not in self.tasks:
            self.task_manage(task)
        self.__dict__["task"] = task

    @abc.abstractmethod
    def compute_scheduler_tags(
        self,
    ) -> tuple[ttags.TaskProvidedTags, wtags.TaskRequiredTags]:
        """Invoke compute_scheduler_tags."""

    def build_expected_tag_candidates(
        self, **kwargs: set[str]
    ) -> tuple[ttags.TaskProvidedTags, wtags.TaskRequiredTags]:
        """
        Build sets of candidate provided and required tags.

        Extra tags can be added by providing sets of tags to arguments named
        ``{provide|require}_{provenance}``
        """
        expected_provided = ttags.TaskProvidedTags()
        for name, tags in kwargs.items():
            if name.startswith("provide_"):
                expected_provided.add(
                    ttags.ProvenanceProvided(name.removeprefix("provide_")),
                    tags,
                )

        expected_required = wtags.TaskRequiredTags()
        for name, tags in kwargs.items():
            if name.startswith("require_"):
                expected_required.add(
                    wtags.ProvenanceRequired(name.removeprefix("require_")),
                    tags,
                )

        return expected_provided, expected_required

    def assertComputesSchedulerTags(self, **kwargs: set[str]) -> None:
        """
        Compute scheduler tags and check that they match the expected values.

        To avoid having to repeat common tags for all tests, this generates two
        expected tag sets with basic values, and lets the user add to them.

        You can specify extra tags to add in arguments named
        ``{provide|require}_{provenance}`` with tag sets to add to the given
        expected tag set.
        """
        expected_provided, expected_required = (
            self.build_expected_tag_candidates(**kwargs)
        )
        provided, required = self.compute_scheduler_tags()
        assert provided.sets == expected_provided.sets
        assert required.sets == expected_required.sets


class ExternalTaskTestCase[TaskClass: BaseExternalTask[Any, Any]](
    TaskTestCase[TaskClass]
):
    """Base test case for external (worker and signing) tasks."""

    #: Field resolver for the test
    resolver: MockFieldResolver

    @override
    def setUp(self) -> None:
        super().setUp()
        self.resolver = MockFieldResolver()

    @override
    def task_setup(self, task: TaskClass) -> None:
        super().task_setup(task)
        # The following values are set by the Worker
        task.work_request_id = 147
        task.workspace_name = "testing"
        task.working_directory = Path(self.create_temporary_directory())

    @override
    def task_create(self, **kwargs: Any) -> TaskClass:
        task_data = self.task_data.copy()
        task_data.update(kwargs)
        task = self.task_class(task_data)
        return task

    def get_subject(self) -> str | None:
        """Compute the subject for this task."""
        assert self.resolver is not None
        self.resolver.resolve_inputs(self.task, stage=inputs.Stage.PENDING)
        return self.task.get_subject()

    def get_configuration_context(self) -> str | None:
        """Compute the configuration context for this task."""
        assert self.resolver is not None
        self.resolver.resolve_inputs(self.task, stage=inputs.Stage.PENDING)
        return self.task.get_configuration_context()

    @override
    def compute_scheduler_tags(
        self, task: TaskClass | None = None
    ) -> tuple[ttags.TaskProvidedTags, wtags.TaskRequiredTags]:
        task = task or self.task
        self.resolver.resolve_inputs(task, stage=inputs.Stage.PENDING)
        return task.compute_scheduler_tags()

    # TODO: this currently cannot use self.task because it cannot be typed to
    # return the task's dynamic data type. Passing self.task is shorter to type
    # than casting or passing the dynamic data type
    def compute_dynamic_data[DTD: BaseDynamicTaskData](
        self, task: BaseExternalTask[Any, DTD]
    ) -> DTD:
        """Compute dynamic data for this task."""
        self.resolver.resolve_inputs(task)
        return task.compute_dynamic_data()

    def mock_debusine(self, task: TaskClass | None = None) -> mock.MagicMock:
        """Create a Debusine mock and configure self.task for it. Return it."""
        task = task or self.task
        debusine_mock = mock.create_autospec(spec=Debusine)
        assert isinstance(debusine_mock, mock.MagicMock)
        # Worker would set the server
        task.configure_server_access(debusine_mock)
        return debusine_mock


class TaskHelperMixin[TaskClass: BaseExternalTask[Any, Any]](
    ExternalTaskTestCase[TaskClass]
):
    """Helper mixin for Task tests."""

    def configure_task(
        self,
        task_data: dict[str, Any] | None = None,
        override: dict[str, Any] | None = None,
        remove: list[str] | None = None,
    ) -> TaskClass:
        """
        Run self.task.configure(task_data) with different inputs.

        Copy self.SAMPLE_TASK_DATA or task_data and modify it, call
        self.task.configure(modified_SAMPLE_TASK_DATA).

        :param task_data: if provided use this as a base dictionary. If None,
          use self.SAMPLE_TASK_DATA.
        :param override: dictionary with the keys to modify (only root keys,
          not recursive search) by its values. Modify or add them.
        :param remove: list of keys to remove.
        """
        if task_data is None:  # pragma: no cover
            task_data = self.SAMPLE_TASK_DATA.copy()
        if override:
            for key, value in override.items():
                task_data[key] = value
        if remove:
            for key in remove:
                task_data.pop(key, None)
        self.task_data = task_data

        # Reset the cached task if one was already built
        self.task_reset_current()
        # Call self.task to instantiate it, as some test methods call
        # configure_task as a way of instantiating a task with the given
        # arguments and see if it breaks
        return self.task

    def can_run_on(self, worker_metadata: dict[str, Any]) -> bool:
        """Invoke can_run_on."""
        self.resolver.resolve_inputs(self.task)
        return self.task.can_run_on(worker_metadata)

    def set_dynamic_data(self) -> None:
        """
        Compute and store the task's dynamic data.

        ``resolver`` must be set up in advance with enough information to
        allow computing dynamic data.
        """
        self.task.dynamic_data = self.compute_dynamic_data(self.task)

    def write_os_release(
        self, filename: Path, codename: str = "bookworm", version: int = 12
    ) -> dict[str, str]:
        """
        Write in directory / self.task._OS_RELEASE_FILE a release file.

        :return: dictionary with keys to values.
        """
        data = {
            "ID": "debian",
            "NAME": "Debian GNU/Linux",
            "PRETTY_NAME": f"Debian GNU/Linux {version} ({codename})",
            "VERSION": f"{version} ({codename})",
            "VERSION_CODENAME": codename,
            "VERSION_ID": str(version),
        }

        if version >= 13:
            del data["VERSION"]
            del data["VERSION_ID"]

        filename.write_text(
            "\n".join(f"{k}={shlex.quote(v)}" for k, v in data.items()) + "\n"
        )
        return data


class ExternalTaskHelperMixin[ExternalTaskClass: BaseExternalTask[Any, Any]](
    TaskHelperMixin[ExternalTaskClass]
):
    """Helper mixin for testing Tasks that run on external workers."""

    task_class: type[ExternalTaskClass]
    task: ExternalTaskClass

    def mock_image_download(
        self,
        debusine_mock: mock.MagicMock,
        create_image: bool = False,
        system_image: bool = False,
    ) -> ArtifactResponse:
        """
        Configure a fake ImageCache path and fake_system_*_artifact.

        The tarball itself isn't faked, unless create_image is set.

        If system_image is set, a system-image is faked instead of a tarball.
        """
        if system_image:
            artifact_response = self.fake_system_image_artifact()
        else:
            artifact_response = self.fake_system_tarball_artifact()
        debusine_mock.artifact_get.return_value = artifact_response

        def download_tarball(
            artifact_id: int,  # noqa: ARG001, U100
            filename: str,  # noqa: ARG001, U100
            destination: Path,
        ) -> None:
            with tarfile.open(destination, "w:xz") as f:
                tarinfo = tarfile.TarInfo("./")
                tarinfo.type = tarfile.DIRTYPE
                f.addfile(tarinfo)
                tarinfo = tarfile.TarInfo("./sbin")
                tarinfo.type = tarfile.DIRTYPE
                f.addfile(tarinfo)
                tarinfo = tarfile.TarInfo("./sbin/init")
                tarinfo.size = 3
                f.addfile(tarinfo, BytesIO(b"123"))

        def download_image(
            artifact_id: int,  # noqa: ARG001, U100
            filename: str,  # noqa: ARG001, U100
            destination: Path,
        ) -> None:
            with destination.open("wb") as f:
                f.write(b"qcow!")

        self.image_cache_path = Path(
            mkdtemp(prefix="debusine-testsuite-images-")
        )
        cast(TestCase, self).addCleanup(rmtree, self.image_cache_path)

        img_cache_patcher = mock.patch(
            "debusine.tasks.executors.images.ImageCache.image_cache_path",
            self.image_cache_path,
        )
        img_cache_patcher.start()
        cast(TestCase, self).addCleanup(img_cache_patcher.stop)

        patcher = mock.patch(
            "debusine.tasks.executors.images.ImageCache"
            "._download_image_artifact"
        )
        self.download_image_artifact_mock = patcher.start()
        cast(TestCase, self).addCleanup(patcher.stop)

        if create_image:
            if system_image:
                self.download_image_artifact_mock.side_effect = download_image
            else:
                self.download_image_artifact_mock.side_effect = download_tarball

        return artifact_response

    def patch_prepare_executor_instance(self) -> mock.MagicMock:
        """
        Patch self.task._prepare_executor_instance(), return its mock.

        Side effect of self.task._prepare_executor_instance(): set
        self.task.executor and self.task.executor_instance.
        """
        patcher = mock.patch.object(
            self.task, "_prepare_executor_instance", autospec=True
        )
        mocked = patcher.start()

        def mock_executor_executor_instance() -> None:
            self.task.executor = mock.MagicMock(spec=ExecutorInterface)
            self.task.executor_instance = mock.MagicMock(spec=UnshareInstance)
            self.task.executor_instance.run.return_value.returncode = 0

        mocked.side_effect = mock_executor_executor_instance
        cast(TestCase, self).addCleanup(patcher.stop)
        return mocked

    def fake_system_tarball_artifact(self) -> ArtifactResponse:
        """Create a fake ArtifactResponse for a debian:system-tarball."""
        return create_artifact_response(
            id=42,
            workspace="Testing",
            category=ArtifactCategory.SYSTEM_TARBALL,
            created_at=dt.datetime(2024, 1, 1, 0, 0, 0, tzinfo=dt.UTC),
            data={
                "architecture": "amd64",
                "codename": "bookworm",
                "filename": "system.tar.xz",
                "vendor": "debian",
                "pkglist": {},
            },
            download_tar_gz_url="https://example.com/download-42/",
            files=FilesResponseType(
                {
                    "system.tar.xz": FileResponse(
                        size=4001,
                        checksums={"sha256": "abc123"},
                        type="file",
                        url="https://example.com/download-system.tar.xz",
                    ),
                }
            ),
            files_to_upload=[],
        )

    def fake_system_image_artifact(self) -> ArtifactResponse:
        """Create a fake ArtifactResponse for a debian:system-image."""
        return create_artifact_response(
            id=42,
            workspace="Testing",
            category=ArtifactCategory.SYSTEM_IMAGE,
            created_at=dt.datetime(2024, 1, 1, 0, 0, 0, tzinfo=dt.UTC),
            data={
                "architecture": "amd64",
                "codename": "bookworm",
                "filename": "image.qcow2",
                "vendor": "debian",
                "image_format": "qcow2",
                "filesystem": "ext4",
                "size": 2 * 1024 * 1024 * 1024,
                "boot_mechanism": "efi",
                "pkglist": {},
            },
            download_tar_gz_url="https://example.com/download-42/",
            files=FilesResponseType(
                {
                    "image.qcow2": FileResponse(
                        size=4001,
                        checksums={"sha256": "abc123"},
                        type="file",
                        url="https://example.com/download-image.qcow2",
                    ),
                }
            ),
            files_to_upload=[],
        )

    def fake_debian_source_package_artifact(self) -> ArtifactResponse:
        """Create a fake ArtifactResponse for a debian:source-package."""
        return create_artifact_response(
            id=6,
            workspace="Testing",
            category=ArtifactCategory.SOURCE_PACKAGE,
            created_at=dt.datetime(2024, 1, 1, 0, 0, 0, tzinfo=dt.UTC),
            data={
                "name": "foo",
                "version": "1.0-1",
                "type": "dpkg",
                "dsc_fields": {},
            },
            download_tar_gz_url="https://example.com/download-6/",
            files=FilesResponseType(
                {
                    "foo_1.0-1.dsc": FileResponse(
                        size=10,
                        checksums={"sha256": "abc123"},
                        type="file",
                        url="https://example.com/foo_1.0-1.dsc",
                    ),
                }
            ),
            files_to_upload=[],
        )

    def fake_debian_binary_package_artifact(self) -> ArtifactResponse:
        """Create a fake ArtifactResponse for a debian:binary-package."""
        return create_artifact_response(
            id=7,
            workspace="Testing",
            category=ArtifactCategory.BINARY_PACKAGE,
            created_at=dt.datetime(2024, 1, 1, 0, 0, 0, tzinfo=dt.UTC),
            data={
                "srcpkg_name": "foo",
                "srcpkg_version": "1.0-1",
                "type": "dpkg",
                "dsc_fields": {},
                "dsc_control_files": [],
            },
            download_tar_gz_url="https://example.com/download-7/",
            files=FilesResponseType(
                {
                    "foo_1.0-1_all.deb": FileResponse(
                        size=10,
                        checksums={"sha256": "abc123"},
                        type="file",
                        url="https://example.com/foo_1.0-1_all.deb",
                    ),
                }
            ),
            files_to_upload=[],
        )

    def fake_debian_upload_artifact(
        self,
        source_package_name: str = "foo",
        architecture: str = "source",
        names: list[str] | None = None,
    ) -> ArtifactResponse:
        """Create a fake ArtifactResponse for a debian:upload."""
        if names is None:
            names = [
                "{source_package_name}_1.0.orig.tar.xz",
                "{source_package_name}_1.0-1.debian.tar.xz",
                "{source_package_name}_1.0-1.dsc",
            ]
        return create_artifact_response(
            id=8,
            workspace="Testing",
            category=ArtifactCategory.UPLOAD,
            created_at=dt.datetime(2024, 1, 1, 0, 0, 0, tzinfo=dt.UTC),
            data={
                "type": "dpkg",
                "changes_fields": {
                    "Architecture": architecture,
                    "Files": [{"name": name} for name in names],
                },
            },
            download_tar_gz_url="https://example.com/download-8/",
            files=FilesResponseType(
                {
                    name: FileResponse(
                        size=len(name),
                        checksums={
                            "sha256": hashlib.sha256(name.encode()).hexdigest()
                        },
                        type="file",
                        url=f"https://example.com/{name}",
                    )
                    for name in [
                        f"{source_package_name}_1.0-1_{architecture}.changes",
                        *names,
                    ]
                }
            ),
            files_to_upload=[],
        )


class SampleBaseExternalTask[
    TD: BaseTaskData,
    DTD: BaseDynamicTaskData = BaseDynamicTaskData,
](
    BaseExternalTask[TD, DTD],
    metaclass=abc.ABCMeta,
):
    """Common test implementation of BaseExternalTask methods."""

    @override
    def get_subject(self) -> str | None:
        raise NotImplementedError()

    @override
    def compute_dynamic_data(self) -> DTD:
        """Resolve artifact lookups for this task."""
        return self.dynamic_task_data_type()

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


class SampleBaseTaskWithExecutor[
    TDE: BaseTaskDataWithExecutor,
    DTDE: BaseDynamicTaskDataWithExecutor,
](SampleBaseExternalTask[TDE, DTDE], BaseTaskWithExecutor[TDE, DTDE]):
    """Common test implementation of BaseExternalTask methods."""
