# Copyright 2022-2023 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 datetime
import tarfile
from io import BytesIO
from pathlib import Path
from shutil import rmtree
from tempfile import mkdtemp
from typing import Any, Generic, Protocol, TypeVar, cast
from unittest import TestCase, mock
from unittest.mock import MagicMock

try:
    import pydantic.v1 as pydantic
except ImportError:
    import pydantic as pydantic  # type: ignore

from debusine.artifacts.models import ArtifactCategory, CollectionCategory
from debusine.client.debusine import Debusine
from debusine.client.models import (
    ArtifactResponse,
    FileResponse,
    FilesResponseType,
)
from debusine.tasks import BaseExternalTask, BaseTask
from debusine.tasks.executors import ExecutorInterface, UnshareInstance
from debusine.tasks.models import LookupMultiple, LookupSingle
from debusine.tasks.server import TaskDatabaseInterface
from debusine.utils import extract_generic_type_arguments


TaskClass = TypeVar("TaskClass", bound=BaseTask[Any, Any])
ExternalTaskClass = TypeVar(
    "ExternalTaskClass", bound=BaseExternalTask[Any, Any]
)


class TaskTestProtocol(Protocol, Generic[TaskClass]):
    """Attributes that tests using :class:`TaskHelperMixin` must provide."""

    SAMPLE_TASK_DATA: dict[str, Any]
    task_class: type[TaskClass]
    task: TaskClass


class FakeTaskDatabase(TaskDatabaseInterface):
    """Fake implementation of database interaction in worker tests."""

    def __init__(
        self,
        single_lookups: (
            dict[tuple[LookupSingle, str | None], int] | None
        ) = None,
        multiple_lookups: (
            dict[tuple[LookupMultiple, str | None], list[int]] | None
        ) = None,
    ) -> None:
        """Construct a :py:class:`FakeTaskDatabase`."""
        self.single_lookups = single_lookups or {}
        self.multiple_lookups = multiple_lookups or {}

    def lookup_single_artifact(
        self,
        lookup: LookupSingle,
        default_category: CollectionCategory | None = None,  # noqa: U100
    ) -> int:
        """Pretend to look up a single artifact."""
        return self.single_lookups[(lookup, default_category)]

    def lookup_multiple_artifacts(
        self,
        lookup: LookupMultiple,
        default_category: CollectionCategory | None = None,  # noqa: U100
    ) -> list[int]:
        """Pretend to look up multiple artifacts."""
        if not lookup.__root__:
            return []
        return self.multiple_lookups[(lookup, default_category)]


class TaskHelperMixin(Generic[TaskClass]):
    """Helper mixin for Task tests."""

    task_class: type[TaskClass]
    task: TaskClass

    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, TaskHelperMixin)

    def configure_task(
        self: TaskTestProtocol[TaskClass],
        task_data: dict[str, Any] | None = None,
        override: dict[str, Any] | None = None,
        remove: list[str] | None = None,
    ):
        """
        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 = self.task_class(task_data)


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

    task_class: type[ExternalTaskClass]
    task: ExternalTaskClass

    def mock_debusine(self: TaskTestProtocol[ExternalTaskClass]) -> MagicMock:
        """Create a Debusine mock and configure self.task for it. Return it."""
        debusine_mock = mock.create_autospec(spec=Debusine)

        # Worker would set the server
        self.task.configure_server_access(debusine_mock)

        return debusine_mock

    def mock_image_download(
        self,
        debusine_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: U100
            filename: str,  # noqa: U100
            destination: Path,
        ):
            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: U100
            filename: str,  # noqa: U100
            destination: Path,
        ):
            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) -> 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():
            self.task.executor = MagicMock(spec=ExecutorInterface)
            self.task.executor_instance = MagicMock(spec=UnshareInstance)

        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 ArtifactResponse(
            id=42,
            workspace="Testing",
            category=ArtifactCategory.SYSTEM_TARBALL,
            created_at=datetime.datetime(
                2024, 1, 1, 0, 0, 0, tzinfo=datetime.UTC
            ),
            data={
                "architecture": "amd64",
                "codename": "bookworm",
                "filename": "system.tar.xz",
                "vendor": "debian",
                "pkglist": {},
            },
            download_tar_gz_url=pydantic.parse_obj_as(
                pydantic.AnyUrl, "https://example.com/download-42/"
            ),
            files=FilesResponseType(
                {
                    "system.tar.xz": FileResponse(
                        size=4001,
                        checksums={"sha256": "abc123"},
                        type="file",
                        url=pydantic.parse_obj_as(
                            pydantic.AnyUrl,
                            "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 ArtifactResponse(
            id=42,
            workspace="Testing",
            category=ArtifactCategory.SYSTEM_IMAGE,
            created_at=datetime.datetime(
                2024, 1, 1, 0, 0, 0, tzinfo=datetime.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=pydantic.parse_obj_as(
                pydantic.AnyUrl, "https://example.com/download-42/"
            ),
            files=FilesResponseType(
                {
                    "image.qcow2": FileResponse(
                        size=4001,
                        checksums={"sha256": "abc123"},
                        type="file",
                        url=pydantic.parse_obj_as(
                            pydantic.AnyUrl,
                            "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 ArtifactResponse(
            id=6,
            workspace="Testing",
            category=ArtifactCategory.SOURCE_PACKAGE,
            created_at=datetime.datetime(
                2024, 1, 1, 0, 0, 0, tzinfo=datetime.UTC
            ),
            data={
                "name": "foo",
                "version": "1.0-1",
                "type": "dpkg",
                "dsc_fields": {},
            },
            download_tar_gz_url=pydantic.parse_obj_as(
                pydantic.AnyUrl, "https://example.com/download-6/"
            ),
            files=FilesResponseType(
                {
                    "foo_1.0-1.dsc": FileResponse(
                        size=10,
                        checksums={"sha256": "abc123"},
                        type="file",
                        url=pydantic.parse_obj_as(
                            pydantic.AnyUrl,
                            "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 ArtifactResponse(
            id=7,
            workspace="Testing",
            category=ArtifactCategory.BINARY_PACKAGE,
            created_at=datetime.datetime(
                2024, 1, 1, 0, 0, 0, tzinfo=datetime.UTC
            ),
            data={
                "srcpkg_name": "foo",
                "srcpkg_version": "1.0-1",
                "type": "dpkg",
                "dsc_fields": {},
                "dsc_control_files": [],
            },
            download_tar_gz_url=pydantic.parse_obj_as(
                pydantic.AnyUrl, "https://example.com/download-7/"
            ),
            files=FilesResponseType(
                {
                    "foo_1.0-1_all.deb": FileResponse(
                        size=10,
                        checksums={"sha256": "abc123"},
                        type="file",
                        url=pydantic.parse_obj_as(
                            pydantic.AnyUrl,
                            "https://example.com/foo_1.0-1_all.deb",
                        ),
                    ),
                }
            ),
            files_to_upload=[],
        )
