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

"""Unit tests for the piuparts task support on the worker."""
import datetime
import lzma
import tarfile
from io import BytesIO
from pathlib import Path
from unittest import TestCase, mock
from unittest.mock import MagicMock, call

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

from debusine.artifacts.models import ArtifactCategory, CollectionCategory
from debusine.client.models import (
    ArtifactResponse,
    FileResponse,
    FilesResponseType,
)
from debusine.tasks import TaskConfigError
from debusine.tasks.models import LookupMultiple, PiupartsDynamicData
from debusine.tasks.piuparts import Piuparts
from debusine.tasks.tests.helper_mixin import (
    ExternalTaskHelperMixin,
    FakeTaskDatabase,
)
from debusine.test import TestHelpersMixin


class PiupartsTaskTests(
    TestHelpersMixin, ExternalTaskHelperMixin[Piuparts], TestCase
):
    """Test the Piuparts Task class."""

    SAMPLE_TASK_DATA = {
        "input": {"binary_artifacts": [421]},
        "host_architecture": "amd64",
        "environment": "debian/match:codename=bookworm",
        "base_tgz": 44,
    }

    def setUp(self):  # noqa: D102
        self.task = Piuparts(self.SAMPLE_TASK_DATA)
        self.task.dynamic_data = PiupartsDynamicData(
            environment_id=1, input_binary_artifacts_ids=[421], base_tgz_id=44
        )

    def tearDown(self):
        """Delete directory to avoid ResourceWarning with python -m unittest."""
        if self.task._debug_log_files_directory is not None:
            self.task._debug_log_files_directory.cleanup()

    def test_configure_fails_with_missing_required_data(self):  # noqa: D102
        with self.assertRaises(TaskConfigError):
            self.configure_task(override={"input": {"artifact": 1}})

    def test_configure_environment_is_required(self):
        """Missing required field "environment": raise exception."""
        with self.assertRaises(TaskConfigError):
            self.configure_task(remove=["environment"])

    def test_configure_base_tgz_is_required(self):
        """Missing required field "base_tgz": raise exception."""
        with self.assertRaises(TaskConfigError):
            self.configure_task(remove=["base_tgz"])

    def test_configure_host_architecture_is_required(self):
        """Missing required field "host_architecture": raise exception."""
        with self.assertRaises(TaskConfigError):
            self.configure_task(remove=["host_architecture"])

    def test_compute_dynamic_data(self) -> None:
        """Dynamic data receives relevant artifact IDs."""
        binary_artifacts_lookup = LookupMultiple.parse_obj([421])
        task_db = FakeTaskDatabase(
            single_lookups={
                # environment
                (
                    "debian/match:codename=bookworm:architecture=amd64:"
                    "format=tarball:backend=unshare",
                    CollectionCategory.ENVIRONMENTS,
                ): 1,
                # base_tgz
                (44, CollectionCategory.ENVIRONMENTS): 44,
            },
            multiple_lookups={
                # input.binary_artifacts
                (binary_artifacts_lookup, None): [421]
            },
        )

        self.assertEqual(
            self.task.compute_dynamic_data(task_db),
            PiupartsDynamicData(
                environment_id=1,
                input_binary_artifacts_ids=[421],
                base_tgz_id=44,
            ),
        )

    def test_cmdline_as_root(self):
        """_cmdline_as_root() return True."""
        self.assertTrue(self.task._cmdline_as_root())

    def patch_fetch_artifact_for_basetgz(
        self,
        *,
        with_dev: bool,
        with_resolvconf_symlink: bool = False,
        codename: str = "bookworm",
        filename: str = "system.tar.xz",
    ) -> MagicMock:
        """Patch self.task.fetch_artifact(), pretend to download basetgz."""
        patcher = mock.patch.object(self.task, "fetch_artifact", autospec=True)
        mocked = patcher.start()

        def _create_basetgz_file(
            artifact_id: int,
            destination: Path,
            default_category: CollectionCategory | None = None,  # noqa: U100
        ):
            with tarfile.open(destination / filename, "w:xz") as tar:
                tarinfo = tarfile.TarInfo(name="./foo")
                tarinfo.size = 12
                tar.addfile(tarinfo, BytesIO(b"A tar member"))
                if with_resolvconf_symlink:
                    tarinfo = tarfile.TarInfo(name="./etc")
                    tarinfo.type = tarfile.DIRTYPE
                    tar.addfile(tarinfo)
                    tarinfo = tarfile.TarInfo(name="./etc/resolv.conf")
                    tarinfo.type = tarfile.SYMTYPE
                    tarinfo.linkname = "../run/systemd/resolve/stub-resolv.conf"
                    tar.addfile(tarinfo)

            data = {
                "with_dev": with_dev,
                "codename": codename,
                "filename": filename,
            }

            file_model = FileResponse(
                size=100,
                checksums={"sha256": "not-used"},
                type="file",
                url=pydantic.parse_obj_as(pydantic.AnyUrl, "https://not-used"),
            )

            files = FilesResponseType({"system.tar.xz": file_model})

            return ArtifactResponse(
                id=artifact_id,
                workspace="System",
                category=ArtifactCategory.SYSTEM_TARBALL,
                created_at=datetime.datetime.now(),
                data=data,
                download_tar_gz_url=pydantic.parse_obj_as(
                    pydantic.AnyUrl,
                    f"http://localhost/artifact/{artifact_id}?archive=tar.gz",
                ),
                files_to_upload=[],
                files=files,
            )

        mocked.side_effect = _create_basetgz_file

        self.addCleanup(patcher.stop)

        return mocked

    def test_configure_for_execution_from_artifact_id_error_no_debs(self):
        """configure_for_execution() no .deb-s: return False."""
        download_directory = self.create_temporary_directory()
        (file1 := download_directory / "file1.dsc").write_text("")
        (file2 := download_directory / "file2.changes").write_text("")

        self.assertFalse(self.task.configure_for_execution(download_directory))

        files = sorted([str(file1), str(file2)])
        log_file_contents = (
            Path(self.task._debug_log_files_directory.name)
            / "configure_for_execution.log"
        ).read_text()
        self.assertEqual(
            log_file_contents,
            f"There must be at least one *.deb file. "
            f"Current files: {files}\n",
        )

    def test_configure_for_execution_prepare_executor_base_tgz(self):
        """configure_for_execution() prepare the executor and base tgz."""
        download_directory = self.create_temporary_directory()
        (download_directory / "package1.deb").write_bytes(b"pkg1")

        self.patch_prepare_executor_instance()

        with mock.patch.object(
            self.task, "_prepare_base_tgz", autospec=True
        ) as prepare_base_tgz_mocked:
            self.assertTrue(
                self.task.configure_for_execution(download_directory)
            )

        self.task.executor_instance.mkdir.assert_called_once_with(
            Path(self.task._piuparts_tmpdir), parents=True
        )
        run_called_with = [
            call(["apt-get", "update"], run_as_root=True, check=True),
            call(
                ["apt-get", "--yes", "install", "piuparts", "mmdebstrap"],
                run_as_root=True,
                check=True,
            ),
        ]
        self.task.executor_instance.run.assert_has_calls(run_called_with)

        prepare_base_tgz_mocked.assert_called_with(download_directory)

    def test_prepare_base_tgz_raise_assertion_error(self):
        """_prepare_base_tgz raise AssertionError: executor_instance is None."""
        msg = r"^self\.executor_instance cannot be None$"
        self.assertRaisesRegex(
            AssertionError,
            msg,
            self.task._prepare_base_tgz,
            self.create_temporary_directory(),
        )

    def test_prepare_base_tgz_download_base_tgz_no_processing(self):
        """
        _prepare_base_tgz does not do any artifact processing (with_dev=False).

        * Download the artifact to the correct directory
        * Set self.task._base_tar to it
        """
        codename = "bullseye"
        filename = "system-image.tar.xz"
        fetch_artifact_mocked = self.patch_fetch_artifact_for_basetgz(
            with_dev=False, codename=codename, filename=filename
        )

        # Use mock's side effect to set self.task.executor_instance
        self.patch_prepare_executor_instance()
        self.task._prepare_executor_instance()

        download_directory = self.create_temporary_directory()

        self.task._prepare_base_tgz(download_directory)

        destination_dir = download_directory / "base_tar"

        fetch_artifact_mocked.assert_called_with(
            self.task.data.base_tgz, destination_dir
        )
        self.assertEqual(self.task._base_tar, destination_dir / filename)
        self.assertEqual(self.task._distribution_codename, codename)

    def test_prepare_base_tgz_download_base_tgz_remove_dev_files(self):
        r"""_prepare_base_tgz download the artifact and remove /dev/\*."""
        # Use mock's side effect to set self.task.executor_instance
        self.patch_prepare_executor_instance()
        self.task._prepare_executor_instance()

        download_directory = self.create_temporary_directory()
        destination_dir = download_directory / "base_tar"

        fetch_artifact_mocked = self.patch_fetch_artifact_for_basetgz(
            with_dev=True
        )

        system_tar = destination_dir / "system.tar"

        self.task._prepare_base_tgz(download_directory)

        fetch_artifact_mocked.assert_called_with(
            self.task.data.base_tgz, destination_dir
        )

        # Downloaded system.tar.xz, then processed it and renamed to system.tar
        self.assertEqual(self.task._base_tar, system_tar)

        self.assertEqual(
            self.task.executor_instance.run.call_args[0][0],
            [
                "mmtarfilter",
                "--path-exclude=/dev/*",
            ],
        )

        reading_from_file = self.task.executor_instance.run.call_args[1][
            "stdin"
        ]
        writing_to_file = self.task.executor_instance.run.call_args[1]["stdout"]

        self.assertEqual(reading_from_file.mode, "rb")
        self.assertEqual(
            reading_from_file.name, str(destination_dir / "system.tar.xz")
        )

        self.assertEqual(writing_to_file.mode, "wb")
        self.assertEqual(
            writing_to_file.name, str(destination_dir / "system.tar")
        )

        self.assertEqual(
            list(destination_dir.iterdir()), [destination_dir / "system.tar"]
        )

    def test_prepare_base_tgz_download_base_tgz_resolveconf(self):
        r"""_prepare_base_tgz download the artifact and install resolv.conf."""
        # Use mock's side effect to set self.task.executor_instance
        self.patch_prepare_executor_instance()
        self.task._prepare_executor_instance()

        download_directory = self.create_temporary_directory()
        destination_dir = download_directory / "base_tar"

        fetch_artifact_mocked = self.patch_fetch_artifact_for_basetgz(
            with_dev=False, with_resolvconf_symlink=True
        )

        def fake_mmtarfilter(args, stdin, stdout):  # noqa: U100
            """Fake for run() that decompresses the tar."""
            stdout.write(lzma.decompress(stdin.read()))

        self.task.executor_instance.run.side_effect = fake_mmtarfilter

        self.task._prepare_base_tgz(download_directory)

        fetch_artifact_mocked.assert_called_with(
            self.task.data.base_tgz, destination_dir
        )

        self.assertEqual(
            self.task.executor_instance.run.call_args[0][0],
            [
                "mmtarfilter",
                "--path-exclude=/etc/resolv.conf",
            ],
        )

        # Downloaded system.tar.xz, then processed it and renamed to system.tar
        system_tar = destination_dir / "system.tar"
        self.assertEqual(self.task._base_tar, system_tar)

        with tarfile.open(system_tar, "r") as tar:
            self.assertTrue(tar.getmember("./etc/resolv.conf").isreg())

    def test_execute(self):
        """Test full (mocked) execution."""
        self.configure_task(override={"input": {'binary_artifacts': [1, 2]}})
        self.task.dynamic_data = PiupartsDynamicData(
            environment_id=1, input_binary_artifacts_ids=[1, 2], base_tgz_id=44
        )
        self.patch_prepare_executor_instance()
        download_directory = self.create_temporary_directory()

        with mock.patch.object(
            self.task,
            "fetch_artifact",
            autospec=True,
            return_value=True,
        ) as mocked:
            self.assertTrue(self.task.fetch_input(download_directory))
            mocked.assert_any_call(1, download_directory)
            mocked.assert_any_call(2, download_directory)

        (file2 := download_directory / "file2.deb").write_text("")
        (file1 := download_directory / "file1.deb").write_text("")
        (file3 := download_directory / "makedev_2.3.1-97_all.deb").write_text(
            ""
        )
        distribution_codename = "trixie"
        self.patch_fetch_artifact_for_basetgz(
            with_dev=True, codename=distribution_codename
        )

        self.assertTrue(self.task.configure_for_execution(download_directory))

        self.assertEqual(self.task._deb_files, [file1, file2, file3])
        self.assertEqual(
            self.task._cmdline(),
            [
                "/usr/sbin/piuparts",
                f"--distribution={distribution_codename}",
                "--allow-database",
                "--warn-on-leftovers-after-purge",
                "--tmpdir=/var/cache/piuparts/tmp",
                f"--basetgz={self.task._base_tar}",
                str(file1),
                str(file2),
                str(file3),
            ],
        )

        self.assertTrue(
            self.task.upload_artifacts(
                download_directory, execution_success=True
            )
        )
