#!/usr/bin/env python3

#    linux-deployment-scripts
#    Copyright (C) 2024  VM-Experiments
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <https://www.gnu.org/licenses/>.

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):
    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'")

    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 ubuntu_rootfs_prepare(
    download_link, download_name, packages, identifier, scripts, package_manager
):
    rootfs_downloader(download_link, download_name)
    generic_rootfs_prepare(packages, identifier, scripts, package_manager)


def debian_rootfs_prepare(
    download_link, download_name, packages, identifier, scripts, package_manager
):
    command(f"sudo debootstrap --arch amd64 {download_name} {chroot} {download_link}")
    generic_rootfs_prepare(packages, identifier, scripts, package_manager)


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


# 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"]
                identifier = data["identifier"]
                scripts = data["scripts"]
                package_manager = data["package_manager"]

                match rootfs_type:
                    case "download":
                        ubuntu_rootfs_prepare(
                            download_link,
                            download_name,
                            packages,
                            identifier,
                            scripts,
                            package_manager,
                        )
                    case "debootstrap":
                        debian_rootfs_prepare(
                            download_link,
                            download_name,
                            packages,
                            identifier,
                            scripts,
                            package_manager,
                        )

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


main()