#!/usr/bin/env python3

# This file is part of VM-Experiments.
# Licensed under the GPL-3.0-or-later. See LICENSE for details.

import subprocess
import sys
import os
import json
import shutil

downloads = "downloads/"
output = "output/"
chroot = "chroot/"
resources = "resources/"
installer_packages = "linux-image-amd64 grub-pc parted tar dosfstools util-linux e2fsprogs arch-install-scripts"
installer_amd64 = "http://deb.debian.org/debian/"
installer_amd64_name = "bookworm"
not_verbose = True
clean_tarballs = False


def command(command):
    result = subprocess.run(command, capture_output=not_verbose, text=True, shell=True)
    return result


# chroot associated functions
def prepare_chroot():
    command(f"sudo mount --bind /dev '{chroot}/dev'")
    command(f"sudo mount --bind /proc '{chroot}/proc'")
    command(f"sudo mount --bind /sys '{chroot}/sys'")
    command(f"sudo cp /etc/resolv.conf '{chroot}/etc/'")


def do_chroot_command(command_input):
    command(f"sudo chroot '{chroot}' /bin/bash -c '{command_input}'")


def close_chroot():
    command(f"sudo umount '{chroot}/dev'")
    command(f"sudo umount '{chroot}/proc'")
    command(f"sudo umount '{chroot}/sys'")


# handle rootfs image from internet
def rootfs_downloader(link, filename):
    os.chdir(downloads)
    command(f"wget -nc {link}")
    command(f"sudo tar -xzf {filename} -C '../{chroot}'")
    os.chdir("..")


def rootfs_creator(filename):
    os.chdir(output)
    command(f"sudo tar -czf {filename} -C '../{chroot}' .")
    os.chdir("..")


# packaging rootfs
def rootfs_package(identifier, scripts, extra_files):
    infile = f"{identifier}.tar.gz"
    outfile = f"{identifier}.iso"

    close_chroot()
    command(f"sudo rm -rf '{chroot}'")
    command(f"sudo mkdir '{chroot}'")

    command(
        f"sudo debootstrap --arch amd64 {installer_amd64_name} {chroot} {installer_amd64}"
    )
    command(f"sudo cp '{infile}' '{chroot}/rootfs.tar.gz'")
    if clean_tarballs:
        command(f"sudo rm '{infile}'")

    prepare_chroot()
    do_chroot_command("apt update")
    do_chroot_command(
        f"DEBIAN_FRONTEND=noninteractive apt install -y {installer_packages}"
    )
    do_chroot_command("mkdir -p /etc/systemd/system/getty@tty1.service.d/")
    do_chroot_command("touch /etc/systemd/system/getty@tty1.service.d/override.conf")
    do_chroot_command(
        'echo -e "[Service]\nExecStart=\nExecStart=/sbin/agetty --noclear --autologin root %I \\$TERM" | tee /etc/systemd/system/getty@tty1.service.d/override.conf'
    )
    do_chroot_command("systemctl daemon-reload")
    do_chroot_command("systemctl enable getty@tty1.service")
    do_chroot_command("echo 'localmachine' > /etc/hostname")
    close_chroot()

    command(f"sudo cp '{resources}/grub.cfg' '{chroot}/boot/grub/grub.cfg'")
    for script in scripts:
        result = script.split("/")[-1]
        command(f"sudo cp '{script}' '{chroot}/root/{result}'")
        command(f"echo './{result}' | sudo tee -a '{chroot}/root/.profile'")
        command(f"echo './{result}' | sudo tee -a '{chroot}/root/.bashrc'")

    for extra_file in extra_files:
        command(f"sudo cp '{extra_file}' '{chroot}'")

    command(f"echo 'umount -R /mnt' | sudo tee -a '{chroot}/root/.profile'")
    command(f"echo 'umount -R /mnt' | sudo tee -a '{chroot}/root/.bashrc'")
    command(f"echo 'reboot' | sudo tee -a '{chroot}/root/.profile'")
    command(f"echo 'reboot' | sudo tee -a '{chroot}/root/.bashrc'")

    if shutil.which("grub-mkrescue"):
        command(f"sudo grub-mkrescue -o '{outfile}' '{chroot}'")
    elif shutil.which("grub2-mkrescue"):
        command(f"sudo grub2-mkrescue -o '{outfile}' '{chroot}'")

    command(f"sudo chown $(logname):$(logname) {outfile}")


# distro specific
def download_rootfs_prepare(
    download_link,
    download_name,
    packages,
    repositories,
    replace_repositories,
    keyrings,
    identifier,
    scripts,
    package_manager,
    extra_files,
):
    rootfs_downloader(download_link, download_name)
    apt_rootfs_prepare(
        repositories, replace_repositories, packages, identifier, scripts, extra_files
    )


def debootstrap_rootfs_prepare(
    download_link,
    download_name,
    packages,
    repositories,
    replace_repositories,
    keyrings,
    identifier,
    scripts,
    package_manager,
    extra_files,
):
    command(f"sudo debootstrap --arch amd64 {download_name} {chroot} {download_link}")
    apt_rootfs_prepare(
        repositories,
        replace_repositories,
        keyrings,
        packages,
        identifier,
        scripts,
        extra_files,
    )


def apt_rootfs_prepare(
    repositories,
    replace_repositories,
    keyrings,
    packages,
    identifier,
    scripts,
    extra_files,
):
    prepare_chroot()

    if replace_repositories:
        overwritten = False
    else:
        overwritten = True

    for repository in repositories:
        if overwritten:
            do_chroot_command(f"echo '{repository}' >>/etc/apt/sources.list")
        else:
            do_chroot_command(f"echo '{repository}' >/etc/apt/sources.list")
            overwritten + True

    for download, name in keyrings:
        do_chroot_command(f"wget {download}")
        do_chroot_command(f"dpkg -i {name}")
        do_chroot_command(f"rm {name}")

    do_chroot_command("apt update -y")
    do_chroot_command("apt full-upgrade -y")
    do_chroot_command(f"DEBIAN_FRONTEND=noninteractive apt install -y {packages}")
    do_chroot_command("echo 'localmachine' > /etc/hostname")
    close_chroot()
    rootfs_creator(f"{identifier}.tar.gz")
    rootfs_package(f"{output}{identifier}", scripts, extra_files)


# main
def main():
    global not_verbose
    global clean_tarballs

    for index, arg in enumerate(sys.argv):
        match arg:
            case "-v":
                not_verbose = False
            case "--verbose":
                not_verbose = False
            case "-ct":
                clean_tarballs = True
            case "--clean-tarballs":
                clean_tarballs = True
            case "-h":
                print("")
                print("Help for Linux Deployment Scripts:")
                print("")
                print("-h / --help - Shows this menu.")
                print(
                    "-v / --verbose - Displays the outputs of the different commands."
                )
                print(
                    "-ct / --clean-tarballs - Delete the tar.gz leftovers from the iso generation."
                )
                print("")
                exit()
            case "--help":
                clean_tarballs = True
            case _:
                if index == 0:
                    continue

                command(f"mkdir -p '{downloads}' '{output}'")
                close_chroot()
                command(f"sudo rm -rf '{chroot}'")
                command(f"sudo mkdir '{chroot}'")

                filename = arg
                data = json.load(open(filename, "r"))

                rootfs_type = data["rootfs_type"]
                download_link = data["download_link"]
                download_name = data["download_name"]
                packages = data["packages"]
                repositories = data["repositories"]
                replace_repositories = bool(data["replace_repositories"])
                keyrings = data["keyrings"]
                identifier = data["identifier"]
                scripts = data["scripts"]
                package_manager = data["package_manager"]
                extra_files = data["extra_files"]

                match rootfs_type:
                    case "download":
                        download_rootfs_prepare(
                            download_link,
                            download_name,
                            packages,
                            repositories,
                            replace_repositories,
                            keyrings,
                            identifier,
                            scripts,
                            package_manager,
                            extra_files,
                        )
                    case "debootstrap":
                        debootstrap_rootfs_prepare(
                            download_link,
                            download_name,
                            packages,
                            repositories,
                            replace_repositories,
                            keyrings,
                            identifier,
                            scripts,
                            package_manager,
                            extra_files,
                        )

    close_chroot()
    command(f"sudo rm -rf '{chroot}'")


main()