#!/usr/bin/env python3
import enum
import subprocess
import sys
import os
import time

BOLD = '\033[1m'
ITALIC = "\033[3m"

UNDERLINE = '\033[4m'
STRIKETHROUGH = '\033[9m'

C_RESET = '\033[0m'

# Basic Colors
BLACK = "\033[30m"
RED = "\033[31m"
GREEN = "\033[32m"
YELLOW = "\033[33m"
BLUE = "\033[34m"
MAGENTA = "\033[35m"
CYAN = "\033[36m"
WHITE = "\033[37m"
RESET = "\033[39m"
GRAY = "\033[90m"

# Light Colors
LIGHT_RED = "\033[91m"
LIGHT_GREEN = "\033[92m"
LIGHT_YELLOW = "\033[93m"
LIGHT_BLUE = "\033[94m"
LIGHT_MAGENTA = "\033[95m"
LIGHT_CYAN = "\033[96m"
LIGHT_WHITE = "\033[97m"

# Background Colors
BACK_BLACK = "\033[40m"
BACK_RED = "\033[41m"
BACK_GREEN = "\033[42m"
BACK_YELLOW = "\033[43m"
BACK_BLUE = "\033[44m"
BACK_MAGENTA = "\033[45m"
BACK_CYAN = "\033[46m"
BACK_WHITE = "\033[47m"
BACK_RESET = "\033[49m"

# Light Background Colors
BACK_LIGHT_RED = "\033[101m"
BACK_LIGHT_GREEN = "\033[102m"
BACK_LIGHT_YELLOW = "\033[103m"
BACK_LIGHT_BLUE = "\033[104m"
BACK_LIGHT_MAGENTA = "\033[105m"
BACK_LIGHT_CYAN = "\033[106m"
BACK_LIGHT_WHITE = "\033[107m"

VERSION = "1.2.1"


class Log(enum.Enum):
    ERROR = 1
    WARNING = 12

    BOOTSTRAP = 0

    SUCCESS = 21
    CLEANUP = 22
    PACKAGE = 23
    INITRAMFS = 24


def check_root_permissions():
    return os.geteuid() == 0


def handle_command_error(returncode: int):
    if returncode != 0:
        log(Log.ERROR, f"Something went wrong while configuring the initramfs ...")
        exit(1)


def delete_last_line():
    sys.stdout.write("\x1b[1A")
    sys.stdout.write("\x1b[2K")


def log(log_type: Log, message: str):
    match log_type:
        case Log.ERROR:
            print(f"{BACK_LIGHT_RED}  ERROR  {BACK_RESET}  {message}{C_RESET}")
        case Log.WARNING:
            print(f"{BACK_YELLOW}   WARNING   {BACK_RESET}  {message}{C_RESET}")

        case Log.BOOTSTRAP:
            print(f"{BACK_GREEN}  BOOTSTRAP  {BACK_RESET}  {message}{C_RESET}")
        case Log.CLEANUP:
            print(f"{BACK_MAGENTA}   CLEANUP   {BACK_RESET}  {message}{C_RESET}")
        case Log.PACKAGE:
            print(f"{BACK_CYAN}   PACKAGE   {BACK_RESET}  {message}{C_RESET}")
        case Log.SUCCESS:
            print(f"{BACK_GREEN}   SUCCESS   {BACK_RESET}  {message}{C_RESET}")
        case Log.INITRAMFS:
            print(f"{BACK_GREEN}  INITRAMFS  {BACK_RESET}  {message}{C_RESET}")


class Args:
    def __init__(self):
        self.args = sys.argv[1:]

        self.ignore_root = False
        self.ignore_bootstrap = False

        self.install_dir = os.getcwd()
        self.data_dir = ""

        self.parse_arguments()

    def parse_arguments(self):
        iter_args = iter(self.args)

        for arg in iter_args:
            if arg in ("-t", "--target"):
                try:
                    port = next(iter_args)
                    self.install_dir = str(port)
                except StopIteration:
                    return None
                except ValueError:
                    return None

        return 0


# -- Command messages (help, about) --
def help_message():
    print(f"""{BOLD}{CYAN}{UNDERLINE}[StrawberryOS] os-bootstrap v{VERSION}{C_RESET}\n\
{GREEN}{BOLD}Usage:{RESET} {WHITE}os-bootstrap {CYAN}[command] {RED}[<options>]{C_RESET}\n\n\
{MAGENTA}{BOLD}Commands:{C_RESET}
    {CYAN}{BOLD}help:{C_RESET} Prints this message
    {CYAN}{BOLD}about:{C_RESET} About os-bootstrap
    {CYAN}{BOLD}auth:{C_RESET} Authenticate with your Strawberry ID
    {CYAN}{BOLD}chroot <dir>:{C_RESET} chroot into a directory

    {CYAN}{BOLD}install <packages>:{C_RESET} Installs StrawberryOS and the specified package features
     {BOLD}↳ {MAGENTA}Options:{C_RESET}
            {CYAN}{BOLD}-t, --target <dir>{C_RESET}  Folder where StrawberryOS will be installed  {GREEN}{BOLD}[default: $PWD]{C_RESET}

    {CYAN}{BOLD}configure:{C_RESET} Configure StrawberryOS
     {BOLD}↳ {MAGENTA}Options:{C_RESET}
            {CYAN}{BOLD}-???, --???{C_RESET}             ???        

    {CYAN}{BOLD}initramfs <subcmd>:{C_RESET} Install & configure the Initramfs
     {BOLD}↳ {MAGENTA}Subcommands:{C_RESET}
            {CYAN}{BOLD}install{C_RESET}  Installs the required initramfs tools{C_RESET}
             {BOLD}↳ {MAGENTA}Options:{C_RESET}
                {CYAN}{BOLD}-t, --target <dir>{C_RESET}   Root directory for installing initramfs-tools  {GREEN}{BOLD}[default: $PWD]{C_RESET}
                {CYAN}{BOLD}-d, --data <device>{C_RESET}  Block device for configuring your data directory
    
""")


def about_message():
    print(f"""{BOLD}{CYAN}{UNDERLINE}os-bootstrap v{VERSION}{C_RESET}\n\
{GREEN}{BOLD}Utility for installing a fresh copy of StrawberryOS{C_RESET}
- {RED}StrawberryOS{C_RESET}, a project by {BLUE}Strawberry Foundations {C_RESET}
""")


def install(cmd_args: Args):
    if not check_root_permissions() and not cmd_args.ignore_root:
        log(Log.ERROR, "Requires root permissions")
        exit(1)

    package_list = {
        "init":
            "dbus,dbus-bin,dbus-daemon,dbus-session-bus-common,dbus-system-bus-common,dbus-user-session,libpam-systemd,\
apt-utils,apt-transport-https,ca-certificates,bash,bzip2,initramfs-tools-core,initramfs-tools,linux-image-amd64,\
busybox-static,network-manager,wget,curl",

        "base": "zstd grub-efi kbd locales locales-all bash-completion sudo",
        "base-dev": "git gcc make g++ build-essential linux-headers-amd64",
        "utils": "neofetch htop btop",
        "python3": "python3 python3-dev python3-pip python-is-python3"
    }

    packages = ["base"]

    i = 0
    while i < len(cmd_args.args):
        if cmd_args.args[i] in ("-t", "--target"):
            try:
                install_dir = str(cmd_args.args.pop(i + 1))
                cmd_args.install_dir = install_dir
                cmd_args.args.pop(i)
            except IndexError:
                log(Log.ERROR, "Please specify a path!")
                exit(1)
            except ValueError:
                log(Log.ERROR, "Invalid value!")
                exit(1)
        else:
            i += 1

    for package in cmd_args.args[1:]:
        if package not in package_list:
            log(Log.ERROR, f"Package {CYAN}'{package}'{C_RESET} is not available")
            exit(1)

    packages.extend(cmd_args.args[1:])

    path = os.path.abspath(cmd_args.install_dir)

    if not os.path.exists(path):
        log(Log.ERROR, f"Path {CYAN}'{path}'{C_RESET} not found!{C_RESET}")
        exit(1)

    for i in reversed(range(0, 4)):
        try:
            log(Log.BOOTSTRAP, f"Starting StrawberryOS installation in {i}s{C_RESET}")
            time.sleep(1)
            delete_last_line()
        except KeyboardInterrupt:
            log(Log.WARNING, f"Canceling StrawberryOS installation {C_RESET}")
            exit(0)

    log(Log.BOOTSTRAP, f"Starting StrawberryOS installation ...{C_RESET}")
    print(f"               {GRAY}-->{C_RESET} Bootstrapping to {CYAN}{path}{RESET} with following packages: {GRAY}{', '.join(packages)}{RESET} ... {C_RESET}\n")

    time.sleep(1)

    if not cmd_args.ignore_bootstrap:
        try:
            command = subprocess.run(
                f"""/usr/sbin/debootstrap \
                --include={package_list["init"]} trixie \
                {path} https://deb.debian.org/debian""", shell=True
            )

            if command.returncode != 0:
                log(Log.ERROR, f"Something went wrong while installing StrawberryOS ...")
                exit(1)

            os.rmdir(path + "/home")

        except Exception as err:
            log(Log.ERROR, f"{err}")
            exit(1)

    cmd = f"bwrap --bind {path} / --dev /dev --bind /sys /sys --bind /proc /proc --bind /tmp /tmp"

    log(Log.BOOTSTRAP, f"Installing additional packages ...{C_RESET}")

    i = 1
    for package in packages:
        log(Log.PACKAGE, f"Installing & configuring package {i} of {len(packages)} ({package}) ...")

        command = subprocess.run(cmd + " apt install -y " + package_list.get(package), shell=True)

        if command.returncode != 0:
            log(Log.ERROR, f"Failed to install package {package} ...")

        i += 1
        print("")

    log(Log.SUCCESS, f"Finished package installation ...{C_RESET}")

    log(Log.CLEANUP, f"Cleaning up system ...{C_RESET}")
    subprocess.run(cmd + " apt update", shell=True)
    subprocess.run(cmd + " apt upgrade -y", shell=True)
    subprocess.run(cmd + " apt clean all", shell=True)
    subprocess.run(cmd + " apt autoclean", shell=True)
    subprocess.run(cmd + " apt autoremove -y", shell=True)

    log(Log.PACKAGE, f"Installing StrawberryOS utils ...")

    subprocess.run(
        cmd + "wget https://raw.githubusercontent.com/Strawberry-Foundations/sbos-scripts/main/update-utils -O /usr/local/bin/update-utils",
        shell=True
    )

    subprocess.run(cmd + "chmod a+x /usr/local/bin/update-utils", shell=True)
    subprocess.run(cmd + "update-utils", shell=True)

    log(Log.SUCCESS, f"Finished bootstrap of base system{C_RESET}")


def initramfs(cmd_args: Args):
    if len(cmd_args.args) <= 1:
        help_message()
        exit(0)

    if not check_root_permissions() and not cmd_args.ignore_root:
        log(Log.ERROR, f"Requires root permissions")
        exit(1)

    match cmd_args.args[1]:
        case "install":
            i = 0
            while i < len(cmd_args.args):
                if cmd_args.args[i] in ("-d", "--data"):
                    try:
                        install_dir = str(cmd_args.args.pop(i + 1))
                        cmd_args.data_dir = install_dir
                        cmd_args.args.pop(i)
                    except IndexError:
                        log(Log.ERROR, f"Please specify a valid block device (e.g. -d /dev/sda2)")
                        exit(1)
                    except ValueError:
                        log(Log.ERROR, "Invalid value!")
                        exit(1)

                elif cmd_args.args[i] in ("-t", "--target"):
                    try:
                        install_dir = str(cmd_args.args.pop(i + 1))
                        cmd_args.install_dir = install_dir
                        cmd_args.args.pop(i)
                    except IndexError:
                        log(Log.ERROR, "Please specify a path!")
                        exit(1)
                    except ValueError:
                        log(Log.ERROR, "Invalid value!")
                        exit(1)
                else:
                    i += 1

            if cmd_args.data_dir == "":
                log(Log.ERROR, f"Please specify a valid block device (e.g. -d /dev/sda2)")
                exit(1)

            path = os.path.abspath(cmd_args.install_dir)

            if not os.path.exists(path):
                log(Log.ERROR, f"Path {CYAN}'{path}'{C_RESET} not found!{C_RESET}")
                exit(1)

            cmd = f"bwrap --bind {path} / --dev /dev --bind /sys /sys --bind /proc /proc --bind /tmp /tmp"

            # Download initramfs tools
            log(Log.INITRAMFS, "Downloading initramfs tools ...")

            command = subprocess.run(
                "wget https://github.com/Strawberry-Foundations/sbos-scripts/archive/refs/heads/main.tar.gz",
                shell=True
            )
            handle_command_error(command.returncode)

            # Extract initramfs tools
            log(Log.INITRAMFS, "Extracting initramfs tools ...")

            command = subprocess.run("tar xfz main.tar.gz", shell=True)
            handle_command_error(command.returncode)

            # Patch initramfs init script
            log(Log.INITRAMFS, "Patching initramfs init script ...")

            with open(f"{os.getcwd()}/sbos-scripts-main/overlay-init", "r") as _initramfs_script:
                initramfs_init = _initramfs_script.read()

            patched_initramfs_init = initramfs_init.replace("/dev/data_partition", cmd_args.data_dir)

            with open(f"{os.getcwd()}/sbos-scripts-main/overlay-init", "w") as _initramfs_script:
                _initramfs_script.write(patched_initramfs_init)


            # Install initramfs tools/scripts
            log(Log.INITRAMFS, f"Installing initramfs-tools to {path} ...")

            command = subprocess.run(f"mv sbos-scripts-main/overlay-hook {path}/etc/initramfs-tools/hooks/overlay", shell=True)
            handle_command_error(command.returncode)

            command = subprocess.run(
                f"mv sbos-scripts-main/overlay-init {path}/etc/initramfs-tools/scripts/init-bottom/overlay",
                shell=True
            )

            handle_command_error(command.returncode)

            with open(f"{os.getcwd()}/sbos-scripts-main/regen-initramfs", "w") as _initramfs_script:
                _initramfs_script.write("update-initramfs -u")

            command = subprocess.run(f"mv sbos-scripts-main/regen-initramfs {path}/usr/local/bin/regen-initramfs", shell=True)
            handle_command_error(command.returncode)

            command = subprocess.run(f"chmod a+x {path}/usr/local/bin/regen-initramfs", shell=True)
            handle_command_error(command.returncode)

            command = subprocess.run(f"chmod a+x {path}/etc/initramfs-tools/hooks/overlay", shell=True)
            handle_command_error(command.returncode)

            command = subprocess.run(f"chmod a+x {path}/etc/initramfs-tools/scripts/init-bottom/overlay", shell=True)
            handle_command_error(command.returncode)

            command = subprocess.run(f"chmod 775 {path}/etc/initramfs-tools/hooks/overlay", shell=True)
            handle_command_error(command.returncode)

            command = subprocess.run(f"chmod 775 {path}/etc/initramfs-tools/scripts/init-bottom/overlay", shell=True)
            handle_command_error(command.returncode)

            log(Log.INITRAMFS, f"Updating initramfs")
            command = subprocess.run(cmd + " regen-initramfs", shell=True)
            handle_command_error(command.returncode)

            log(Log.CLEANUP, f"Cleaning up ...{C_RESET}")
            command = subprocess.run(f"rm -rf sbos-scripts-main", shell=True)
            handle_command_error(command.returncode)

            command = subprocess.run(f"rm -rf main.tar* ", shell=True)
            handle_command_error(command.returncode)

            log(Log.SUCCESS, f"Finished initramfs configuration ")


def chroot(cmd_args: Args):
    if not check_root_permissions() and not cmd_args.ignore_root:
        log(Log.ERROR, f"Requires root permissions")
        exit(1)

    path = os.path.abspath(cmd_args.args[1])

    if not os.path.exists(path):
        log(Log.ERROR, f"Path {CYAN}'{path}'{C_RESET} not found!{C_RESET}")
        exit(1)

    print(f"{BOLD}{GREEN}=>{RESET} Binding host system directories ...{C_RESET}")
    subprocess.run(f"mount --bind /dev {path}/dev", shell=True)
    subprocess.run(f"mount --bind /sys {path}/sys", shell=True)
    subprocess.run(f"mount --bind /proc {path}/proc", shell=True)
    print(f"{BOLD}{GREEN}=>{RESET} Entering chroot ...{C_RESET}")
    subprocess.run(f"chroot {path}", shell=True)
    print(f"{BOLD}{GREEN}=>{RESET} Leaving chroot ...{C_RESET}")
    subprocess.run(f"umount {path}/dev", shell=True)
    subprocess.run(f"umount {path}/sys", shell=True)
    subprocess.run(f"umount {path}/proc", shell=True)


args = Args()

if len(args.args) < 1:
    help_message()
    exit(0)

match args.args[0]:
    case "install":
        install(args)
    case "initramfs":
        initramfs(args)
    case "chroot":
        chroot(args)
    case "about":
        about_message()
    case "help" | _:
        help_message()
