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

"""SimpleSystemImageBuild Task, extending the SystemImageBuild ontology."""
import importlib.resources
import re
import subprocess
from pathlib import Path
from typing import Any, Optional

import yaml

from debusine import utils
from debusine.artifacts.local_artifact import DebianSystemImageArtifact
from debusine.tasks.models import DiskImageFormat
from debusine.tasks.systembootstrap import SystemBootstrap
from debusine.tasks.systemimagebuild import (
    SystemImageBuild,
)

try:
    from importlib.resources.abc import Traversable
except ImportError:
    # Python < 3.12
    from importlib.abc import Traversable


class SimpleSystemImageBuild(SystemImageBuild):
    """Implement SimpleSystemImageBuild using debos."""

    CAPTURE_OUTPUT_FILENAME = "debos.txt"

    ARCH_DEB_TO_DPS = {
        "amd64": "X86_64",
        "arm64": "ARM64",
        "armel": "ARM",
        "armhf": "ARM",
        "i386": "X86",
        "mips64el": "MIPS64_LE",
        "ppc64el": "PPC64_LE",
        "riscv64": "RISCV64",
        "s390x": "S390X",
        "alpha": "ALPHA",
        "hppa": "HPPA",
        "hurd-amd64": "X86_64",
        "hurd-i386": "X86",
        "ia64": "IA64",
        "loong64": "LOONGARCH64",
        "powerpc": "PPC",
        "ppc64": "PPC64",
        "x32": "X86",
    }

    # taken from
    # https://uapi-group.org/specifications/specs/discoverable_partitions_specification/
    DPS_TO_GPT_ROOT = {
        "SD_GPT_ROOT_ALPHA": "6523f8ae-3eb1-4e2a-a05a-18b695ae656f",
        "SD_GPT_ROOT_ARC": "d27f46ed-2919-4cb8-bd25-9531f3c16534",
        "SD_GPT_ROOT_ARM": "69dad710-2ce4-4e3c-b16c-21a1d49abed3",
        "SD_GPT_ROOT_ARM64": "b921b045-1df0-41c3-af44-4c6f280d3fae",
        "SD_GPT_ROOT_IA64": "993d8d3d-f80e-4225-855a-9daf8ed7ea97",
        "SD_GPT_ROOT_LOONGARCH64": "77055800-792c-4f94-b39a-98c91b762bb6",
        "SD_GPT_ROOT_MIPS_LE": "37c58c8a-d913-4156-a25f-48b1b64e07f0",
        "SD_GPT_ROOT_MIPS64_LE": "700bda43-7a34-4507-b179-eeb93d7a7ca3",
        "SD_GPT_ROOT_PARISC": "1aacdb3b-5444-4138-bd9e-e5c2239b2346",
        "SD_GPT_ROOT_PPC": "1de3f1ef-fa98-47b5-8dcd-4a860a654d78",
        "SD_GPT_ROOT_PPC64": "912ade1d-a839-4913-8964-a10eee08fbd2",
        "SD_GPT_ROOT_PPC64_LE": "c31c45e6-3f39-412e-80fb-4809c4980599",
        "SD_GPT_ROOT_RISCV32": "60d5a7fe-8e7d-435c-b714-3dd8162144e1",
        "SD_GPT_ROOT_RISCV64": "72ec70a6-cf74-40e6-bd49-4bda08e8f224",
        "SD_GPT_ROOT_S390": "08a7acea-624c-4a20-91e8-6e0fa67d23f9",
        "SD_GPT_ROOT_S390X": "5eead9a9-fe09-4a1e-a1d7-520d00531306",
        "SD_GPT_ROOT_TILEGX": "c50cdd70-3862-4cc3-90e1-809a8c93ee2c",
        "SD_GPT_ROOT_X86": "44479540-f297-41b2-9af7-d131d5f0458a",
        "SD_GPT_ROOT_X86_64": "4f68bce3-e8cd-4db1-96e7-fbcaf984b709",
    }

    ARCH_TO_GRUB_EFI = {
        "amd64": "x86_64-efi",
        "arm64": "arm64-efi",
        "armhf": "arm-efi",
        "i386": "x86_64-efi",  # We assume 64-bit EFI
        "ia64": "ia64-efi",
        "loong64": "loongarch64-efi",
        "riscv32": "riscv32-efi",
        "riscv64": "riscv64-efi",
    }

    def __init__(self, task_data: dict[str, Any]) -> None:
        """Initialize SimpleSystemImageBuild."""
        super().__init__(task_data)

        self._debos_recipe: Optional[Path] = None
        if self.data.disk_image.format == DiskImageFormat.QCOW2:
            self.filename = f"{self.data.disk_image.filename}.qcow2"
        else:
            self.filename = f"{self.data.disk_image.filename}.tar.xz"

    @classmethod
    def analyze_worker(cls):
        """Report metadata for this task on this worker."""
        metadata = super().analyze_worker()

        available_key = cls.prefix_with_task_name("available")
        metadata[available_key] = utils.is_command_available(
            "debos"
        ) and utils.is_command_available("/lib/systemd/systemd-resolved")
        arch_key = cls.prefix_with_task_name("architecture")
        metadata[arch_key] = (
            subprocess.run(
                ["dpkg-architecture", "-qDEB_HOST_ARCH"], capture_output=True
            )
            .stdout.decode()
            .strip()
        )

        return metadata

    def can_run_on(self, worker_metadata: dict[str, Any]) -> bool:
        """Check if the specified worker can run the task."""
        if not super().can_run_on(worker_metadata):
            return False

        available_key = self.prefix_with_task_name("available")
        arch_key = self.prefix_with_task_name("architecture")
        request_arch = self.data.bootstrap_options.architecture
        host_arch = worker_metadata.get(arch_key)
        # TODO move the architecture comparison to a general place in Debusin
        return worker_metadata.get(available_key, False) and (
            request_arch == host_arch
            or (host_arch == "amd64" and request_arch == "i386")
            or (host_arch == "arm64" and request_arch in ("armel", "armhf"))
        )

    def _cmdline(self) -> list[str]:
        """
        Return debos command line.

        Use configuration of self.data.
        """
        cmd = ["debos", "--verbose", str(self._debos_recipe)]

        return cmd

    def fetch_input(self, destination: Path) -> bool:  # noqa: U100
        """Do nothing: no artifacts need to be downloaded."""
        return True

    def build_debootstrap_actions(
        self, download_dir: Path
    ) -> list[dict[str, Any]]:
        """Build the debootstrap actions for the recipe."""
        suite = self.data.bootstrap_repositories[0].suite

        bootstrap: dict[str, Any] = {
            "action": "debootstrap",
            "description": "Run debootstrap",
            "mirror": self.data.bootstrap_repositories[0].mirror,
            "suite": suite,
        }

        if components := self.data.bootstrap_repositories[0].components:
            bootstrap["components"] = components

        if (variant := self.data.bootstrap_options.variant) is not None:
            bootstrap["variant"] = variant

        if self.data.bootstrap_repositories[0].keyring:
            bootstrap["keyring-file"] = SystemBootstrap._download_key(
                self.data.bootstrap_repositories[0],
                keyring_directory=download_dir,
            )

        keyring_package = self.data.bootstrap_repositories[0].keyring_package
        if keyring_package:
            bootstrap["keyring-package"] = keyring_package

        # Following debootstrap, created non-usr-merged images for <= stretch
        # and <= bookworm for the buildd variant
        if suite in ("jessie", "stretch"):
            bootstrap["merged-usr"] = False
        if variant == "buildd" and suite in ("buster", "bullseye", "bookworm"):
            bootstrap["merged-usr"] = False

        return [bootstrap]

    def build_filesystem_actions(self) -> list[dict[str, Any]]:
        """Build the image and filesystem actions."""
        architecture = self.data.bootstrap_options.architecture
        suite = self.data.bootstrap_repositories[0].suite
        efi_size_mb = 256
        image_size = (
            self.data.disk_image.partitions[0].size + efi_size_mb * 0.001
        )

        root_filesystem = self.data.disk_image.partitions[0].filesystem
        root_features: list[str] = []
        if root_filesystem == "ext4" and suite == "jessie":
            root_features.append("^metadata_csum")

        return [
            {
                "action": "image-partition",
                "imagename": self.image_name,
                "imagesize": f"{image_size}G",
                "partitiontype": "gpt",
                "description": "Create and partition image",
                "mountpoints": [
                    {
                        "mountpoint": "/",
                        "partition": "root",
                    },
                    {
                        "mountpoint": "/efi",
                        "partition": "EFI",
                    },
                ],
                "partitions": [
                    {
                        "name": "EFI",
                        "fs": "vfat",
                        "parttype": ("c12a7328-f81f-11d2-ba4b-00a0c93ec93b"),
                        "start": "0%",
                        "end": f"{efi_size_mb}M",
                        "options": ["x-systemd.automount"],
                    },
                    {
                        "name": "root",
                        "fs": root_filesystem,
                        "parttype": self.DPS_TO_GPT_ROOT[
                            "SD_GPT_ROOT_" + self.ARCH_DEB_TO_DPS[architecture]
                        ],
                        "start": f"{efi_size_mb}M",
                        "end": "100%",
                        "features": root_features,
                    },
                ],
            },
            {
                "action": "filesystem-deploy",
                "setup-kernel-cmdline": True,
                "append-kernel-cmdline": (
                    "console=ttyS0 rootwait rw "
                    "fsck.mode=auto fsck.repair=yes"
                ),
            },
        ]

    def build_incus_agent_actions(self) -> list[dict[str, Any]]:
        """Build the actions to install incus-agent-setup."""
        suite = self.data.bootstrap_repositories[0].suite
        actions = [
            {
                "action": "overlay",
                "source": "overlays/incus-agent",
                "description": "Install incus-agent-setup",
            },
            {
                "action": "run",
                "description": "Make incus-agent-setup executable",
                "command": (
                    "chmod 755 ${ROOTDIR}/lib/systemd/incus-agent-setup"
                ),
            },
        ]
        if suite == "jessie":
            actions.append(
                {
                    "action": "overlay",
                    "source": "overlays/incus-agent-jessie",
                    "description": "Install incus-agent.service for jessie",
                }
            )
        return actions

    def build_init_actions(self) -> list[dict[str, Any]]:
        """Build the actions to install systemd."""
        return [
            {
                "action": "apt",
                "description": "Install init",
                "packages": ["init"],
            }
        ]

    def build_bootloader_actions(self) -> list[dict[str, Any]]:
        """Build the actions to install a bootloader."""
        suite = self.data.bootstrap_repositories[0].suite
        architecture = self.data.bootstrap_options.architecture

        bootloader = self.data.disk_image.bootloader
        if bootloader is None:
            if architecture in ("amd64", "i386", "arm64", "armhf", "riscv64"):
                # In jessie, systemd-boot is not yet available
                if suite == "jessie":
                    bootloader = "grub-efi"
                else:
                    bootloader = "systemd-boot"
            else:
                raise ValueError(f"No default bootloader for {architecture}")

        packages: list[str] = []
        if bootloader == "systemd-boot":
            # In bullseye and below,
            # systemd-boot is part of the systemd package
            if suite not in ("jessie", "stretch", "buster", "bullseye"):
                packages.append("systemd-boot")
        else:
            packages.append(bootloader)

        actions: list[dict[str, Any]] = []

        if packages:
            actions.append(
                {
                    "action": "apt",
                    "description": "Install bootloader",
                    "packages": packages,
                }
            )

        if bootloader == "systemd-boot":
            # bullseye and below require
            # some manual configuration for systemd-boot
            if suite in ("jessie", "stretch", "buster", "bullseye"):
                actions.extend(
                    [
                        {
                            "action": "overlay",
                            "source": "overlays/systemd-boot",
                        },
                        {
                            "action": "run",
                            "description": (
                                "Make systemd-boot .d files executable"
                            ),
                            "command": (
                                "chmod 755 "
                                "${ROOTDIR}/etc/kernel/*/zz-update-systemd-boot"
                            ),
                        },
                        {
                            "action": "run",
                            "description": "Create loader entries directory",
                            "command": "mkdir -p ${ROOTDIR}/efi/$(cat "
                            "${ROOTDIR}/etc/machine-id)",
                        },
                    ]
                )
            actions.append(
                {
                    "action": "run",
                    "description": "Install systemd-boot to /efi",
                    "chroot": True,
                    "command": "bootctl install",
                }
            )
        elif bootloader == "grub-efi":
            grub_target = self.ARCH_TO_GRUB_EFI[architecture]
            actions.extend(
                [
                    {
                        "action": "run",
                        "description": "Install grub-efi to /efi",
                        "chroot": True,
                        "command": (
                            f"grub-install --target={grub_target} "
                            "--efi-directory=/efi --removable"
                        ),
                    },
                    {
                        "action": "run",
                        "description": "Enable serial console kernel cmdline",
                        "command": (
                            "echo GRUB_CMDLINE_LINUX=\"console=ttyS0\" "
                            ">> ${ROOTDIR}/etc/default/grub"
                        ),
                    },
                    {
                        "action": "run",
                        "description": "Write initial grub configuration",
                        "chroot": True,
                        "command": "update-grub",
                    },
                ]
            )
        return actions

    def build_networking_actions(self) -> list[dict[str, Any]]:
        """Build the actions to configure networking."""
        suite = self.data.bootstrap_repositories[0].suite
        actions: list[dict[str, Any]] = [
            {
                "action": "run",
                "description": "Set hostname",
                "command": (
                    "umask 0022 && echo debian > ${ROOTDIR}/etc/hostname"
                ),
            },
        ]
        # In bullseye and below,
        # systemd-resolved is part of the systemd package
        if suite not in ("jessie", "stretch", "buster", "bullseye"):
            actions.append(
                {
                    "action": "apt",
                    "description": "Install systemd-resolved",
                    "packages": ["systemd-resolved"],
                }
            )

        # bullseye and below require
        # manually enabling systemd-resolved
        if suite in ("jessie", "stretch", "buster", "bullseye"):
            actions.append(
                {
                    "action": "run",
                    "description": "Enable systemd-resolved",
                    "command": (
                        "systemctl --root=${ROOTDIR} enable systemd-resolved"
                    ),
                }
            )

        # In jessie we get a slightly different device and configuration
        if suite == "jessie":
            networkd_config = "[Match]\nName=eth*\n[Network]\nDHCP=both"
        else:
            networkd_config = "[Match]\nName=en*\n[Network]\nDHCP=yes"

        actions.extend(
            [
                {
                    "action": "run",
                    "description": "Enable systemd-networkd",
                    "command": "systemctl --root=${ROOTDIR} "
                    "enable systemd-networkd",
                },
                {
                    "action": "run",
                    "description": "Setup network interface",
                    "command": (
                        f"echo '{networkd_config}' "
                        "> ${ROOTDIR}/etc/systemd/network/default.network"
                    ),
                },
            ]
        )
        return actions

    def build_kernel_actions(self) -> list[dict[str, Any]]:
        """Build the action to install the kernel."""
        kernel = self.data.disk_image.kernel_package
        if kernel is None:
            # Only introduced in bullseye
            kernel = "linux-image-generic"
        return [
            {
                "action": "apt",
                "description": "Install kernel",
                "packages": [kernel],
            }
        ]

    def build_apt_actions(self) -> list[dict[str, Any]]:
        """Build the action to configure apt."""
        actions: list[dict[str, Any]] = []

        keyring_package = self.data.bootstrap_repositories[0].keyring_package
        if keyring_package:
            actions.append(
                {
                    "action": "apt",
                    "description": "Install keyring",
                    "packages": [keyring_package],
                }
            )
        extra_packages = self.data.bootstrap_options.extra_packages
        if extra_packages:
            actions.append(
                {
                    "action": "apt",
                    "description": "Install additional packages",
                    "packages": extra_packages,
                }
            )

        return actions

    def build_user_actions(self) -> list[dict[str, Any]]:
        """Build the action to configure users."""
        return [
            {
                "action": "run",
                "description": "Allow login without password",
                "command": "passwd --root=${ROOTDIR} --delete root",
            }
        ]

    def build_query_actions(self) -> list[dict[str, Any]]:
        """Build the actions to query image metadata into the log."""
        return [
            {
                "action": "run",
                "description": "Get os-release",
                "command": "cat ${ROOTDIR}/etc/os-release",
            },
            {
                "action": "run",
                "description": "Get package list",
                "command": "dpkg-query --root ${ROOTDIR} -W",
            },
        ]

    def build_customization_actions(
        self, download_dir: Path
    ) -> list[dict[str, Any]]:
        """Build the actions to run the customization script."""
        cscript = self.data.customization_script
        actions: list[dict[str, Any]] = []

        if cscript is None:
            return actions

        script = {
            "action": "run",
            "description": "Run customization_script",
        }

        # work around setup-testbed not functioning in the debos chroot
        if cscript == "/usr/share/autopkgtest/setup-commands/setup-testbed":
            script["command"] = (
                "env AUTOPKGTEST_KEEP_APT_SOURCES=1 "
                "/usr/share/autopkgtest/setup-commands/setup-testbed "
                "${ROOTDIR}"
            )
        else:
            customization_script = download_dir / "customization_script"
            customization_script.write_text(cscript)

            script["command"] = str(customization_script)
            script["chroot"] = True  # type: ignore

        actions.append(script)
        return actions

    def build_convert_actions(self) -> list[dict[str, Any]]:
        """Build an action to convert the image to the final format."""
        actions: list[dict[str, Any]] = []
        if self.data.disk_image.format == DiskImageFormat.QCOW2:
            actions.append(
                {
                    "action": "run",
                    "description": "Convert image to qcow2",
                    "postprocess": True,
                    "command": (
                        f"qemu-img convert -O qcow2 "
                        f"{self.image_name} {self.filename}"
                    ),
                }
            )
        else:
            actions.append(
                {
                    "action": "run",
                    "description": "tar image",
                    "postprocess": True,
                    "command": (
                        f"tar --create --auto-compress --file {self.filename} "
                        f"{self.image_name}"
                    ),
                }
            )
        return actions

    def configure_for_execution(self, download_dir: Path) -> bool:
        """Prepare debos recipe."""
        self._debos_recipe = download_dir / "debusine.yaml"
        self.image_name = f"{self.data.disk_image.filename}.img"
        architecture = self.data.bootstrap_options.architecture
        recipe: dict[str, Any] = {
            "architecture": architecture,
            "actions": self.build_debootstrap_actions(download_dir)
            + self.build_filesystem_actions()
            + self.build_incus_agent_actions()
            + self.build_init_actions()
            + self.build_bootloader_actions()
            + self.build_networking_actions()
            # kernel install needs to happen after bootctl install
            + self.build_kernel_actions()
            + self.build_apt_actions()
            + self.build_user_actions()
            + self.build_query_actions()
            + self.build_customization_actions(download_dir)
            + self.build_convert_actions(),
        }

        overlays = [
            action["source"].split("/", 1)[1]
            for action in recipe["actions"]
            if action["action"] == "overlay"
        ]
        self.write_overlays(download_dir / "overlays", overlays)

        with open(self._debos_recipe, "w") as rf:
            yaml.safe_dump(recipe, rf)

        return True

    def write_overlays(self, destination: Path, overlays: list[str]):
        """Copy all named overlays into destination."""
        destination.mkdir(exist_ok=True)
        for overlay in overlays:
            self.copy_resource_tree(
                importlib.resources.files(__package__)
                .joinpath("data")
                .joinpath("overlays")
                .joinpath(overlay),
                destination / overlay,
            )

    def copy_resource_tree(self, src: Traversable, dest: Path) -> None:
        """
        Recursive copy a Traversible into a directory.

        Assumes files are tiny.
        Can be replaced with importlib.resources.as_file() in Python >= 3.12.
        """
        if src.is_file():
            with dest.open("wb") as f:
                f.write(src.read_bytes())
        else:
            dest.mkdir(exist_ok=True)
            for member in src.iterdir():
                self.copy_resource_tree(member, dest / member.name)

    def upload_artifacts(self, execute_dir: Path, *, execution_success: bool):
        """Upload generated artifacts."""
        if not self.debusine:
            raise AssertionError("self.debusine not set")

        if not execution_success:
            return

        debos_log = (execute_dir / self.CAPTURE_OUTPUT_FILENAME).read_text()
        pkglist = dict(
            re.findall(
                r"dpkg-query --root \${ROOTDIR} -W \| (.*)\t(.*)", debos_log
            )
        )
        os_release = dict(
            re.findall(
                r"cat \${ROOTDIR}/etc/os-release \| (.*)=(.*)", debos_log
            )
        )

        codename = os_release.get("VERSION_CODENAME", None)
        if codename is None:
            # jessie doesn't provide VERSION_CODENAME
            codename = self.data.bootstrap_repositories[0].suite

        system_file = execute_dir / self.filename
        artifact = DebianSystemImageArtifact.create(
            system_file,
            data={
                "variant": self.data.bootstrap_options.variant,
                "architecture": self.data.bootstrap_options.architecture,
                "vendor": os_release["ID"],
                "codename": codename,
                "pkglist": pkglist,
                "with_dev": True,
                "with_init": True,
                # TODO / XXX: is "mirror" meant to be the first repository?
                # or "mirrors" and list all of them?
                # Or we could duplicate all the bootstrap_repositories...
                "mirror": self.data.bootstrap_repositories[0].mirror,
                "image_format": self.data.disk_image.format,
                "filesystem": self.data.disk_image.partitions[0].filesystem,
                "size": (self.data.disk_image.partitions[0].size * 10**9),
                "boot_mechanism": "efi",
            },
        )

        self.debusine.upload_artifact(
            artifact,
            workspace=self.workspace,
            work_request=self.work_request,
        )
