#!/bin/env python3

# Build a live ISO.
# Run this from an empty working directory,
# or a directory where you have run this before to reuse caches.

import argparse
import subprocess
import pathlib
import tomllib
import hashlib
import datetime
import urllib.request


DISTRIBUTION = "bookworm"

VARIANT_LABEL = {
    "contestant": "contest",
    "training-live": "live",
    "training-installer": "install",
}

VARIANT_BOOT_OPTIONS = {
    "contestant": [
        dict(label="SOI contest", cmdline="boot=live toram=filesystem.squashfs"),
        dict(label="SOI contest, run from drive", cmdline="boot=live"),
        dict(label="SOI contest, fail-safe mode", cmdline="@LB_BOOTAPPEND_LIVE_FAILSAFE@"),
    ],
    "training-live": [
        dict(label="SOI live system, run from RAM", cmdline="boot=live toram=filesystem.squashfs"),
        dict(label="SOI live system, run from drive", cmdline="boot=live"),
        dict(label="SOI live system, fail-safe mode", cmdline="@LB_BOOTAPPEND_LIVE_FAILSAFE@"),
    ],
    "training-installer": [],
}

VARIANT_EXTRA_LB_CONFIG = {
    "training-installer": [
        "--debian-installer", "live",
        # Linux headers are needed for VirtualBox DKMS.
        "--linux-packages", "linux-image linux-headers",
    ],
}

VARIANT_EXTRA_BOOTSTRAP = {
    # Fasttrack keyring is needed for VirtualBox.
    "training-installer": ",fasttrack-archive-keyring",
}

DOWNLOADS = [
    dict(
        name="soi-header.tar.gz",
        url="https://git.soi.ch/SOI/soi-header/archive/19ddcef24eb55bdb5ddb817c1d91bfa04c8cb8dd.tar.gz",
        sha256="38042587982af4e9431aea461e5c345bde358bcc79f0a0eadcf5b3ed77aeb8ab",
    ),
    # From https://soi.ch/wiki/soi-codeblocks/#install-the-soi-project-template
    dict(
        name="soi_template_codeblocks_ubuntu.zip",
        url="https://soi.ch/media/files/soi_template_codeblocks_ubuntu_RzdvSho.zip",
        sha256="3f4cae26bbb0cbdfd4cf9f94bfce14988e395f847948d9b349c12d0b980386e9",
    ),
]

def run(args, check=True, **kwargs):
    print(f"> {' '.join(args)}")
    subprocess.run(args, check=check, **kwargs)

def mkdir(dirname):
    pathlib.Path(dirname).mkdir(parents=True, exist_ok=True)

def edit_file(filename, fn):
    with open(filename) as f:
        content = f.read()
    content = fn(content)
    with open(filename, "w") as f:
        f.write(content)

def sha256sum(filename):
    with open(filename, "rb", buffering=0) as f:
        return hashlib.file_digest(f, "sha256").hexdigest()

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "variant",
        choices=["contestant", "training-live", "training-installer"],
        help="Variant to build.",
    )
    args = parser.parse_args()

    script_dir = pathlib.Path(__file__).parent.resolve()

    with open(script_dir / "config/config.toml", "rb") as f:
        config = tomllib.load(f)

    # Remove files generated by previous build, but keep cache.
    run("lb clean".split())
    run("rm -rf .build config udeb-build".split())

    # Download files.
    mkdir("downloads")
    for download in DOWNLOADS:
        filename = f'downloads/{download["name"]}'
        try:
            if sha256sum(filename) == download["sha256"]:
                continue
        except FileNotFoundError:
            pass
        print(f'> Downloading {download["url"]}')
        urllib.request.urlretrieve(download["url"], filename)
        if sha256sum(filename) != download["sha256"]:
            raise Exception(f"Downloaded file {filename} has wrong hash.")

    # Create base configuration directory.
    run(
        [
            "lb", "config",
            "--clean",
            "--ignore-system-defaults",
            "--mode", "debian",
            "--distribution", DISTRIBUTION,
            "--archive-areas", "main contrib non-free non-free-firmware",
            "--firmware-chroot", "false",
            "--firmware-binary", "false",
            "--apt-recommends", "false",
            # isc-dhcp-client and ifupdown are obsoleted by network-manager,
            # but they still have Priority: important.
            # We need ca-certificates for fetching https packages repos.
            "--debootstrap-options", "--exclude=isc-dhcp-client,isc-dhcp-common,ifupdown --include=ca-certificates" +
                VARIANT_EXTRA_BOOTSTRAP.get(args.variant, ""),
            "--loadlin", "false",
            "--iso-volume", f"SOI {VARIANT_LABEL[args.variant]} @ISOVOLUME_TS@",
            "--bootappend-live", "boot=live toram=filesystem.squashfs",
        ]
        + VARIANT_EXTRA_LB_CONFIG.get(args.variant, [])
    )

    # Add our own configuration on top.
    run(["cp", "-rT", f"{script_dir}/layers/participant", "config"])
    if args.variant != "training-installer":
        run(["cp", "-rT", f"{script_dir}/layers/live", "config"])
    run(["cp", "-rT", f"{script_dir}/layers/{args.variant}", "config"])

    if args.variant == "training-installer":
        # Insert admin password into preseed.
        edit_file("config/includes.installer/preseed.cfg",
            lambda s: s.replace("@install_admin_password@", config["install_admin_password"]))

        # Copy inventory file.
        mkdir("config/includes.binary/install")
        run(["cp", f"{script_dir}/config/installer-inventory.txt", "config/includes.binary/install/inventory.txt"])

        # Insert build date in login screen logo.
        edit_file("config/includes.chroot/usr/local/share/images/login-screen-logo.svg",
            lambda s: s.replace("@date@", datetime.date.today().isoformat()))

        # Build and install custom udeb packages for installer.
        run(["cp", "-rT", f"{script_dir}/installer-udeb", "udeb-build"])
        run("dpkg-buildpackage --build=all".split(), cwd="udeb-build/inventory-hostname")
        mkdir("config/packages.binary")
        run("cp udeb-build/inventory-hostname_0_all.udeb config/packages.binary/".split())

        # Copy the source lists. The installer deletes everything in sources.list.d,
        # so we need to copy them somewhere else and restore them after the install.
        for listpath in pathlib.Path('config/archives').glob('*.list.chroot'):
            run(["cp", str(listpath), f"config/includes.chroot/usr/local/share/target-sources/{listpath.name.removesuffix('.chroot')}"])
        # Insert distribution into source configs.
        for sourcepath in pathlib.Path('config/includes.chroot/usr/local/share/target-sources').glob('*'):
            edit_file(sourcepath, lambda s: s.replace("@DISTRIBUTION@", DISTRIBUTION))
    elif args.variant == "contestant":
        # Insert root password into hook script.
        edit_file("config/hooks/live/2010-contestant.hook.chroot",
            lambda s: s.replace("@contestant_root_password@", config["contestant_root_password"]))

        # Copy authorized_keys.
        mkdir("config/includes.chroot/root/.ssh")
        run(["cp", f"{script_dir}/config/contestant_authorized_keys", "config/includes.chroot/root/.ssh/authorized_keys"])

        edit_file("config/includes.chroot/etc/NetworkManager/system-connections/contest.nmconnection",
            lambda s: s.replace("@wifi_password@", config["contestant_wifi_password"]))

    # Configure boot options.
    grub_boot_options = '\n'.join(
        f'menuentry "{option["label"]}" {{\n'
        f'  linux @KERNEL_LIVE@ {option["cmdline"]}\n'
        f'  initrd @INITRD_LIVE@\n'
        f'}}\n'
        for option in VARIANT_BOOT_OPTIONS[args.variant]
    )
    with open("/usr/share/live/build/bootloaders/grub-pc/grub.cfg") as f:
        grub_cfg = f.read()
    grub_cfg = grub_cfg.replace("@LINUX_LIVE@", grub_boot_options)
    with open("config/bootloaders/grub-pc/grub.cfg", "w") as f:
        f.write(grub_cfg)

    syslinux_boot_options = ''.join(
        f"label live-{i}\n"
        f"\tmenu label {option['label']}\n"
        + (f"\tmenu default\n" if i == 0 else "") +
        f"\tlinux @LINUX@\n"
        f"\tinitrd @INITRD@\n"
        f"\tappend {option['cmdline'].replace('@LB_BOOTAPPEND_LIVE_FAILSAFE@', '@APPEND_LIVE_FAILSAFE@')}\n"
        f"\n"
        for i, option in enumerate(VARIANT_BOOT_OPTIONS[args.variant])
    )
    mkdir("config/bootloaders/syslinux_common")
    with open("config/bootloaders/syslinux_common/live.cfg.in", "w") as f:
        f.write(syslinux_boot_options)

    # Install soi header.
    mkdir("config/includes.chroot/usr/local/include")
    run(["tar", "--overwrite", "-xf", "downloads/soi-header.tar.gz", "-C", "config/includes.chroot/usr/local/include", "--strip-components=2", "soi-header/include/"])

    # Install codeblocks template.
    mkdir("config/includes.chroot/usr/share/codeblocks/templates/wizard")
    run(["unzip", "-o", "downloads/soi_template_codeblocks_ubuntu.zip", "-d", "config/includes.chroot/usr/share/codeblocks/templates/wizard/"])

    # Start the build.
    run("lb build".split())

if __name__ == "__main__":
    main()