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

"""
Task to test Debian binary packages using piuparts.

This task follows the specification laid out here:
https://freexian-team.pages.debian.net/debusine/reference/tasks.html#piuparts-task
"""

import gzip
import re
import shlex
import shutil
import tarfile
from collections.abc import Collection
from io import BytesIO
from pathlib import Path
from tempfile import TemporaryDirectory
from typing import override

from debian.debian_support import Version

import debusine.utils
from debian import debfile
from debusine.artifacts import PiupartsArtifact
from debusine.artifacts.models import (
    ArtifactCategory,
    DebianBinaryPackage,
    DebianPiuparts,
    DebianSystemTarball,
    DebianUpload,
    WorkRequestResults,
    get_source_package_name,
    get_source_package_version,
)
from debusine.client.models import RelationType
from debusine.tasks import (
    BaseTaskWithExecutor,
    ExtraRepositoryMixin,
    RunCommandTask,
    TaskConfigError,
)
from debusine.tasks.executors import ExecutorImageCategory
from debusine.tasks.inputs import EnvironmentInput, MultiInput, Stage
from debusine.tasks.models import BackendType, PiupartsData, PiupartsDynamicData


class Piuparts(
    ExtraRepositoryMixin[PiupartsData, PiupartsDynamicData],
    RunCommandTask[PiupartsData, PiupartsDynamicData],
    BaseTaskWithExecutor[PiupartsData, PiupartsDynamicData],
):
    """Test Debian binary packages using piuparts."""

    TASK_VERSION = 1

    CAPTURE_OUTPUT_FILENAME = "piuparts.txt"

    base_tgz = EnvironmentInput(
        image_category=ExecutorImageCategory.TARBALL,
        set_backend=False,
        try_variant=False,
        categories=[ArtifactCategory.SYSTEM_TARBALL],
        stage=Stage.PENDING,
    )
    environment = EnvironmentInput()
    input_binary_artifacts = MultiInput(
        field="input.binary_artifacts",
        categories=[ArtifactCategory.BINARY_PACKAGE, ArtifactCategory.UPLOAD],
        stage=Stage.PENDING,
    )

    @override
    def post_init(self) -> None:
        super().post_init()

        # *.deb paths and versions. Set by self.configure_for_execution()
        self._deb_files: list[Path] = []
        self._deb_versions: dict[str, str] | None = None

        # Version of piuparts in the executor.
        # Set by self.configure_for_execution()
        self._piuparts_version: Version | None = None

        # --basetgz=... option for piuparts cmdline.
        # Set by self.configure_for_execution()
        self._base_tar: Path | None = None
        self._base_tar_data: DebianSystemTarball | None = None
        self._scripts_dir: Path | None = None

    @override
    def get_subject(self) -> str:
        package_names = set()
        for binary_artifact in self.input_binary_artifacts.artifacts:
            assert isinstance(
                binary_artifact.data,
                (DebianBinaryPackage, DebianUpload),
            )
            package_names.add(get_source_package_name(binary_artifact.data))

        package_names_sorted = sorted(package_names)
        return " ".join(package_names_sorted)

    @override
    def get_configuration_context(self) -> str | None:
        assert isinstance(self.base_tgz.data, DebianSystemTarball)
        return self.base_tgz.data.codename

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

        :subject: space-separated list of source package names from
          ``input.binary_artifact``
        :parameter_summary: same as subject with the version (if one package)
        :runtime_context: ``codename:build_architecture`` where ``codename``
          is from ``base_tgz``
        :configuration_context: ``codename`` from ``base_tgz``
        """
        base_tgz = self.base_tgz
        assert isinstance(base_tgz.data, DebianSystemTarball)

        source_package_names = set()
        source_versions = set()
        for binary_artifact in self.input_binary_artifacts.artifacts:
            assert isinstance(
                binary_artifact.data,
                (DebianBinaryPackage, DebianUpload),
            )
            source_package_names.add(
                get_source_package_name(binary_artifact.data)
            )
            source_versions.add(
                get_source_package_version(binary_artifact.data)
            )

        runtime_context = f"{base_tgz.data.codename}:{self.build_architecture}"

        subject = self.get_subject()
        if len(source_package_names) == 1 and len(source_versions) == 1:
            (source_version,) = source_versions
            (source_package_name,) = source_package_names
            parameter_summary = f"{source_package_name}_{source_version}"
        else:
            parameter_summary = subject

        # TODO: It would probably be better to get the architecture of the
        # input binary artifacts and use that.
        parameter_summary += f" ({self.build_architecture})"

        return PiupartsDynamicData(
            environment_id=self.environment.id,
            extra_repositories=self.extra_repositories,
            input_binary_artifacts_ids=self.input_binary_artifacts.resolved_ids(),
            base_tgz_id=base_tgz.id,
            subject=subject,
            parameter_summary=parameter_summary,
            runtime_context=runtime_context,
            configuration_context=self.get_configuration_context(),
        )

    @override
    def fetch_input(self, destination: Path) -> bool:
        """Populate work directory with user-specified binary artifact(s)."""
        assert self.dynamic_data

        for artifact_id in self.dynamic_data.input_binary_artifacts_ids:
            self.fetch_artifact(artifact_id, destination)
        return True

    @override
    def configure_for_execution(self, download_directory: Path) -> bool:
        r"""
        Find the .deb files for piuparts.

        Prepare executor, install "piuparts" in it and prepare
        debian:system-image (e.g. delete /dev/\* files).

        Set self._deb_files to the relevant files.

        :param download_directory: where to find the \*.deb files (downloaded
           via fetch_input) and where to download the chroot of
           debian:system-image (for piuparts --basetgz).
        :return: True if valid files were found
        """
        # Find the files to test or early exit if no files
        self._deb_files = debusine.utils.find_files_suffixes(
            download_directory, [".deb"]
        )
        # Ensure we've got >=1 .deb file(s)
        # Note: This could be moved into a more generic function shared across
        # tasks.
        if len(self._deb_files) == 0:
            list_of_files = sorted(map(str, download_directory.iterdir()))
            self.append_to_log_file(
                "configure_for_execution.log",
                [
                    f"There must be at least one *.deb file. "
                    f"Current files: {list_of_files}"
                ],
            )
            return False

        deb_versions: dict[str, str] = {}
        for file in self._deb_files:
            pkg = debfile.DebFile(file)
            try:
                control = pkg.control.debcontrol()
                deb_versions[control["Package"]] = control["Version"]
            finally:
                pkg.close()
        self._deb_versions = deb_versions

        # Prepare executor
        self._prepare_executor_instance()

        if self.executor_instance is None:
            raise AssertionError("self.executor_instance cannot be None")

        # Prepare executor_instance to run piuparts
        self.run_executor_command(
            ["apt-get", "update"],
            log_filename="install.log",
            run_as_root=True,
            check=True,
        )

        self.run_executor_command(
            ["apt-get", "--yes", "install", "piuparts"],
            log_filename="install.log",
            run_as_root=True,
            check=True,
        )

        self._check_piuparts_version()

        self._prepare_base_tgz(download_directory)
        self._scripts_dir = download_directory / "scripts"
        self._scripts_dir.mkdir()
        self._prepare_scripts(self._scripts_dir)

        return True

    def _check_piuparts_version(self) -> None:
        """Check that the environment contains a supported piuparts version."""
        # We are root, older piuparts will work
        if self.backend in (BackendType.INCUS_VM, BackendType.QEMU):
            return

        # We are inside a container, and can't create device nodes (#638)
        self._piuparts_version = self.get_package_version("piuparts")
        min_version = Version("1.5")

        if self._piuparts_version < min_version:
            raise TaskConfigError(
                f"The environment contains a version of piuparts "
                f"({self._piuparts_version}) that is too old. "
                f"Container-based backends require piuparts >= "
                f"{min_version}. We recommend a >= trixie environment "
                f"with an older base_tgz."
            )

    def _prepare_base_tgz(self, download_directory: Path) -> None:
        """
        Download the base image for piuparts. Remove /dev/* if needed.

        Set self._base_tar to the file that is used by _cmdline.
        """
        if self.executor_instance is None:
            raise AssertionError("self.executor_instance cannot be None")
        assert self.dynamic_data

        (base_tar_dir := download_directory / "base_tar").mkdir()

        base_tgz_response = self.fetch_artifact(
            self.dynamic_data.base_tgz_id, base_tar_dir
        )
        self._base_tar_data = DebianSystemTarball.model_validate(
            base_tgz_response.data
        )
        self._base_tar = base_tar_dir / self._base_tar_data.filename

        # Is the base_tgz set up for systemd-resolved?
        resolvconf_symlink = False
        # Does /var/lib/dpkg/available exist?
        dpkg_available = False
        with tarfile.open(self._base_tar) as tar:
            first_member = tar.getmembers()[0]
            # How are the members prefixed? Typically ./
            m = re.search(r"^(?:\.?/)?", first_member.name)
            assert m  # This regex would match an empty string
            prefix = m.group(0)

            try:
                if tar.getmember(
                    f"{prefix}etc/resolv.conf"
                ).issym():  # pragma: no cover
                    resolvconf_symlink = True
            except KeyError:
                pass

            try:
                if tar.getmember(
                    f"{prefix}var/lib/dpkg/available"
                ):  # pragma: no cover
                    dpkg_available = True
            except KeyError:
                pass

        if (
            not base_tgz_response.data["with_dev"]
            and not resolvconf_symlink
            and dpkg_available
        ):
            # No processing needs to be done, the image can stay as
            # .tar.xz (or the format retrieved)
            return

        # The tar will be fixed up and self._base_tar pointed at the new file

        filter_dirs: list[Path] = []
        filter_files: list[Path] = []
        create_empty_files: list[Path] = []
        if base_tgz_response.data["with_dev"]:
            filter_dirs.append(Path("/dev"))
        if resolvconf_symlink:
            # Replace a broken symlink with an empty file
            filter_files.append(Path("/etc/resolv.conf"))
            create_empty_files.append(Path("/etc/resolv.conf"))
        if not dpkg_available:
            create_empty_files.append(Path("/var/lib/dpkg/available"))

        processed_tar_gz = self._filter_tar(
            self._base_tar,
            filter_dirs=filter_dirs,
            filter_files=filter_files,
            create_empty_files=create_empty_files,
        )

        # Use the processed image
        self._base_tar = processed_tar_gz

    def _filter_tar(
        self,
        input_: Path,
        filter_dirs: Collection[Path] = (),
        filter_files: Collection[Path] = (),
        create_empty_files: Collection[Path] = (),
    ) -> Path:
        """Process a tarball, filtering contents and creating empty files."""
        processed_tar_name = input_.stem.rsplit(".", 1)[0] + ".tar"
        output = input_.with_name(processed_tar_name)
        if input_ == output:
            raise AssertionError(
                "We generated a filtered filename that matched the input."
            )

        with (
            tarfile.open(input_) as input_tar,
            tarfile.open(output, mode="w") as output_tar,
        ):
            for member in input_tar:
                path = Path("/" + member.name.lstrip('./'))
                if any(path.is_relative_to(dir_) for dir_ in filter_dirs):
                    continue
                if path in filter_files:
                    continue

                if member.isfile():
                    output_tar.addfile(
                        member, input_tar.extractfile(member.name)
                    )
                else:
                    output_tar.addfile(member)

            for path in create_empty_files:
                for parent in reversed(path.parents):
                    if parent == Path("/"):
                        continue
                    try:
                        if output_tar.getmember(f".{parent}").isdir():
                            continue
                        raise AssertionError(
                            f"{parent} exists in the tarball, but not a "
                            "directory."
                        )
                    except KeyError:
                        tarinfo = tarfile.TarInfo(name=f".{parent}")
                        tarinfo.mode = 0o755
                        tarinfo.type = tarfile.DIRTYPE
                        output_tar.addfile(tarinfo)
                tarinfo = tarfile.TarInfo(name=f".{path}")
                tarinfo.mode = 0o644
                tarinfo.size = 0
                output_tar.addfile(tarinfo, BytesIO(b""))

        # Compress the processed tarball for compatibility with piuparts < 1.3
        output_tar_gz = output.with_suffix(".tar.gz")
        with (
            open(output, "rb") as uncompressed,
            gzip.open(output_tar_gz, "wb") as compressed,
        ):
            shutil.copyfileobj(uncompressed, compressed)

        output.unlink()

        return output_tar_gz

    def _prepare_scripts(self, script_directory: Path) -> None:
        """Prepare scripts for piuparts."""
        if not self._base_tar_data:
            raise AssertionError(
                "self._base_tar_data not set, call _prepare_base_tgz first."
            )

        bad_permissions_script = (
            script_directory / "pre_remove_50_find_bad_permissions"
        )
        bad_permissions_script.write_text(
            "#!/bin/sh\n"
            "echo 'Skipping to avoid https://bugs.debian.org/1125784'\n"
        )
        bad_permissions_script.chmod(0o755)

        install_passwd_script = script_directory / "post_setup_00_passwd"
        install_passwd_script.write_text(
            "#!/bin/sh\n"
            "echo 'Installing passwd for https://bugs.debian.org/1125789'\n"
            "if [ ! -e /usr/sbin/useradd ]; then\n"
            "    apt-get -y install passwd\n"
            "fi\n"
        )
        install_passwd_script.chmod(0o755)

        codename = self._base_tar_data.codename
        if self.data.extra_repositories:
            extra_repo_script = (
                script_directory / "post_chroot_unpack_debusine_extra_repos"
            )
            with (
                TemporaryDirectory(
                    prefix="debusine-piuparts-"
                ) as extra_repositories_dir,
                extra_repo_script.open("w") as f,
            ):
                extra_repositories_path = Path(extra_repositories_dir)
                self.write_extra_repository_config(
                    codename=codename, destination=extra_repositories_path
                )

                f.write("#!/bin/sh\nset -eux\n")
                if self.extra_repository_keys:
                    f.write("mkdir -p /etc/apt/keyrings\n")
                for key in self.extra_repository_keys:
                    f.write(
                        f"printf %s {shlex.quote(key.read_text())}"
                        f" > /etc/apt/keyrings/{key.name}\n"
                    )
                for source in (
                    self.extra_repository_sources + self.extra_repository_lists
                ):
                    f.write(
                        f"printf %s {shlex.quote(source.read_text())}"
                        f" > /etc/apt/sources.list.d/{source.name}\n"
                    )
                f.write("echo DONE\n")
            extra_repo_script.chmod(0o755)

    @override
    def _cmdline(self) -> list[str]:
        """Build full piuparts command line."""
        cmdline = ["piuparts"]

        # Use --keep-sources-list (not --distribution=dist-variant /
        # --mirror / --extra-repo): we feed pre-generated tarballs and
        # we don't want piuparts to overwrite the APT configuration.
        # Also don't set --distribution/-d: we don't want to
        # work-around dated distros.conf aliasing (stable/testing/...)
        # and variants (-backports,-updates,/updates,-security/...).
        # Only use -d in the future if we need to enable --scriptsdir
        # and need to set $PIUPARTS_DISTRIBUTION.
        cmdline.append("--keep-sources-list")

        # Common options used in Debian piuparts services/CI
        cmdline.append("--allow-database")
        cmdline.append("--warn-on-leftovers-after-purge")

        # Use pre-generated tarballs (don't debootstrap on the fly)
        cmdline.append(f"--basetgz={self._base_tar}")

        cmdline.append("--scriptsdir=/etc/piuparts/scripts")

        # Configure APT sources, etc.
        cmdline.append(f"--scriptsdir={self._scripts_dir}")

        # Test the binary .deb files
        cmdline.extend(map(str, self._deb_files))

        return cmdline

    @override
    @staticmethod
    def _cmdline_as_root() -> bool:
        """Piuparts must run as root."""
        return True

    @override
    def check_directory_for_consistency_errors(
        self, build_directory: Path
    ) -> list[str]:
        """Return list of errors."""
        if build_directory / self.CAPTURE_OUTPUT_FILENAME not in set(
            build_directory.iterdir()
        ):
            return [f"{self.CAPTURE_OUTPUT_FILENAME} not in {build_directory}"]

        return []

    @override
    def upload_artifacts(
        self,
        execute_directory: Path,
        *,
        execution_result: WorkRequestResults,  # noqa: U100
    ) -> None:
        """Upload a PiupartsArtifact with the files, data, and relationships."""
        assert self.debusine is not None
        assert self._deb_versions is not None
        assert self._piuparts_version is not None
        assert self._base_tar_data is not None

        base_tar_data = self._base_tar_data
        piuparts_artifact = PiupartsArtifact.create(
            piuparts_output=execute_directory / self.CAPTURE_OUTPUT_FILENAME,
            data=DebianPiuparts(
                piuparts_version=str(self._piuparts_version),
                binary_packages=self._deb_versions,
                # It may seem as though this should be the architecture of
                # the binary packages instead (for example, we might be
                # testing `Architecture: all` binary packages on amd64).
                # However, the results of running `piuparts` inherently
                # depend on the build architecture in other ways even if the
                # input artifacts all happen to be architecture-independent.
                architecture=self.data.build_architecture,
                distribution=f"{base_tar_data.vendor}:{base_tar_data.codename}",
            ),
        )

        uploaded_piuparts_artifact = self.debusine.upload_artifact(
            piuparts_artifact,
            workspace=self.workspace_name,
            work_request=self.work_request_id,
        )

        for source_artifact_id in self._source_artifacts_ids:
            self.debusine.relation_create(
                uploaded_piuparts_artifact.id,
                source_artifact_id,
                RelationType.RELATES_TO,
            )
