#!/usr/bin/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

downloads = "downloads/"
output = "output/"
chroot = "chroot/"
resources = "resources/"
installer_packages = "linux-image-generic grub-pc parted tar dosfstools util-linux e2fsprogs arch-install-scripts"
installer_amd64 = (
    "https://cdimage.ubuntu.com/ubuntu-base/noble/daily/current/noble-base-amd64.tar.gz"
)
installer_amd64_name = "noble-base-amd64.tar.gz"


def command(command):
    result = subprocess.run(command, capture_output=False, 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(filename, installer):
    infile = f"{filename}.tar.gz"
    outfile = f"{filename}.iso"

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

    rootfs_downloader(installer_amd64, installer_amd64_name)
    command(f"sudo cp '{infile}' '{chroot}/rootfs.tar.gz'")

    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")
    close_chroot()

    command(f"sudo cp '{resources}/grub.cfg' '{chroot}/boot/grub/grub.cfg'")
    command(f"sudo cp '{resources}/{installer}' '{chroot}/root'")
    command(f"echo './{installer}' | sudo tee -a '{chroot}/root/.profile'")
    command(f"echo './{installer}' | sudo tee -a '{chroot}/root/.bashrc'")
    command(f"sudo grub2-mkrescue -o '{outfile}' '{chroot}'")


# distro specific
def ubuntu_rootfs_prepare(link, filename, packages, outputname, script):
    rootfs_downloader(link, filename)
    prepare_chroot()
    do_chroot_command("apt update -y")
    do_chroot_command("apt full-upgrade -y")
    do_chroot_command(f"DEBIAN_FRONTEND=noninteractive apt install -y {packages}")
    close_chroot()
    rootfs_creator(f"{outputname}.tar.gz")
    rootfs_package(f"{output}{outputname}", script)


def main():
    for index, arg in enumerate(sys.argv):
        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"))

        link = data["link"]
        download_name = data["download_name"]
        packages = data["packages"]
        identifier = data["identifier"]
        script = data["script"]

        ubuntu_rootfs_prepare(link, download_name, packages, identifier, script)

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


main()