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

CONFIG_PATH = os.path.dirname(os.path.abspath(__file__))

for cfg in ["config.json", "default-config.json"]:
    if os.path.isfile(CONFIG_PATH + "/" + cfg):
        CONFIG_PATH = CONFIG_PATH + "/" + cfg
        break

def load_config():
    with open(CONFIG_PATH, "r") as f:
        return json.load(f)

config = load_config()
def detect_primary_disk():
    for disk in os.listdir("/sys/block"):
        if disk.startswith("dm") \
           or disk.startswith("fd") \
           or disk.startswith("ram") \
           or disk.startswith("zram") \
           or disk.startswith("loop"):
            continue
        with open(f"/sys/block/{disk}/removable", "r") as f:
            if f.read().strip() == "1":
                continue
        if "usb" in os.path.realpath(f"/sys/block/{disk}"):
            continue
        prefix=""
        if disk.startswith("mmcblk") or disk.startswith("nvme"):
            prefix="p"
        return disk, prefix
    return None, None

def cttyhack(console="tty1"):
    sys.stdin = open(f"/dev/{console}", "r")
    sys.stdout = open(f"/dev/{console}", "w")
    sys.stderr = open(f"/dev/{console}", "w")

def fallback(e):
    print("=========== Error ===========")
    print(e)
    print("========= Traceback =========")
    print(traceback.format_exc())
    print("=============================")
    print("\x1b[31;1mInstallation failed.\x1b[0m")
    print("Creating a shell for debuging. Good luck :D")
    envs = {
        "PATH": os.environ["PATH"],
        "PS1": "\x1b[32;1m>>>\x1b[0m "}
    subprocess.run(["/bin/bash", "--norc", "--noprofile"], env=envs, check=False)
    quit(1)

def create_parts(disk, prefix):

    format_table = {
        "ext4" : ["-F", "-L", "LABEL"],
        "btrfs": ["-f", "-L", "LABEL"],
        "vfat" : ["-n", "LABEL"]
    }

    def run_parted(args):
        try:
            subprocess.run(["parted", "-s", f"/dev/{disk}"] + args, check=True)
        except Exception as e:
            fallback(e)
    def run_mkfs(fstype, num=1, label="primary"):
        args = format_table.get(fstype, [])
        if "LABEL" in args:
            i = args.index("LABEL")
            args[i] = label
        try:
            subprocess.run([f"mkfs.{fstype}"] + (args or []) + [f"/dev/{disk}{prefix}{num}"], check=True)
        except Exception as e:
            fallback(e)

    def run_mount(num, target):
        try:
            path = f"/target/{target}"
            os.makedirs(path, exist_ok=True)
            subprocess.run(["mount", f"/dev/{disk}{prefix}{num}", path], check=True)
        except Exception as e:
            fallback(e)
    # create mbr
    if os.path.isdir("/sys/firmware/efi"):
        run_parted(["mktable", "gpt"])
    else:
        run_parted(["mktable", "msdos"])

    # create parts
    if os.path.isdir("/sys/firmware/efi"):
        run_parted(["mkpart", "primary", "fat32", "1", "100M"])
        run_parted(["mkpart", "primary", "fat32", "100M", "100%"])
        run_parted(["set", "1", "esp", "on"])
    else:
        run_parted(["mkpart", "primary", "fat32", "1", "100%"])

    # format
    if os.path.isdir("/sys/firmware/efi"):
        run_mkfs("vfat", 1, "EFI")
        run_mkfs(config.get("rootfs_type", "ext4"), 2, "ROOTFS")
    else:
        run_mkfs(config.get("rootfs_type", "ext4"), 1, "ROOTFS")

    # mount
    if os.path.isdir("/sys/firmware/efi"):
        run_mount(2, "/")
        run_mount(1, "/boot/efi")
    else:
        run_mount(1, "/")

def rsync():
    try:
        os.makedirs("/source", exist_ok=True)
        subprocess.run(["mount", "/dev/loop0", "/source"], check=True)
        subprocess.run(
            ["rsync", "--archive", "--info=progress2",
            "--no-inc-recursive", "--human-readable", "--hard-links",
            "--xattrs", "--acls", "-r", "/source/", "/target/"],
            check=True
        )
    except Exception as e:
        fallback(e)

def bind(unbind=False):
    try:
        for d in ["dev", "sys", "proc", "run"]:
            if unbind:
                subprocess.run(
                    ["umount", "-lf", "-R", f"/target/{d}"], check=True
                )
            else:
                subprocess.run(
                    ["mount", "--bind", f"/{d}", f"/target/{d}"], check=True
                )
    except Exception as e:
        fallback(e)

def remove_live():
    cmds = [
        ["apt-get", "purge", "live-boot", "live-config", "live-tools", "--yes"],
        ["apt-get", "autoremove", "--yes"],
        ["update-initramfs", "-u", "-k", "all"]
    ]
    try:
        for cmd in cmds:
            subprocess.run(
                ["chroot", "/target"] + cmd, check=True
            )
    except Exception as e:
        fallback(e)

def install_grub(disk):
    def run_cmd(cmd):
        try:
            subprocess.run(["chroot", "/target"] + cmd, check=True)
        except Exception as e:
            fallback(e)
    # efivarfs
    target="i386-pc"
    if os.path.isdir("/sys/firmware/efi"):
        run_cmd(["mount", "-t", "efivarfs", "efivarfs", "/sys/firmware/efi/efivars"])
        target="x86_64-efi"
    # install grub
    run_cmd(["grub-install", f"/dev/{disk}", f"--target={target}"])
    run_cmd(["grub-install", f"/dev/{disk}", f"--target={target}", "--removable"])
    # update grub
    run_cmd(["grub-mkconfig", "-o", "/boot/grub/grub.cfg"])

def find_uuid(disk, prefix="", num=1):
    sb = subprocess.run(["blkid", "-o", "json", f"/dev/{disk}{prefix}{num}"], stdout=subprocess.PIPE)
    if sb.returncode != 0:
        return None
    data = json.loads(sb.stdout.decode("utf-8"))
    blocks = data.get("blkid") or []
    if not blocks:
        return None
    return blocks[0].get("uuid")


def write_fstab(disk, prefix):
    fstab = open("/target/etc/fstab", "w")
    def fstab_line(num, target, fstype, option):
        device = f"/dev/{disk}{prefix}{num}"
        uuid = find_uuid(disk, prefix, num)
        ret = f"# {device}\n"
        if uuid:
            device = f"UUID={uuid}"
        ret += f"{device}  {target} {fstype} {option}\n"
        return ret

    if os.path.isdir("/sys/firmware/efi"):
        fstab.write(fstab_line(2, "/", config.get("rootfs_type", "ext4"), "defaults,rw 0 1"))
        fstab.write(fstab_line(1, "/boot/efi", "vfat", "defaults,rw 0 0"))
    else:
        fstab.write(fstab_line(1, "/", config.get("rootfs_type", "ext4"), "defaults,rw 0 1"))
    fstab.flush()
    fstab.close()

def write_x11_keyboard():
    os.makedirs("/target/etc/X11/xorg.conf.d", exist_ok=True)
    with open("/target/etc/X11/xorg.conf.d/10-keyboard.conf", "w") as f:
        f.write('Section "InputClass"\n')
        f.write('Identifier "system-keyboard"\n')
        f.write('MatchIsKeyboard "on"\n')
        f.write(f'Option "XkbLayout" "{config["keyboard_layout"]}"\n')
        f.write(f'Option "XkbModel" "{config["keyboard_model"]}"\n')
        f.write(f'Option "XkbVariant" "{config["keyboard_variant"]}"\n')
        f.write('EndSection\n')


def write_language():
    with open("/target/etc/locale.gen", "w") as f:
        for line in config.get("locale_gen", []):
            f.write(f"{line}\n")
    with open("/target/etc/default/locale", "w") as f:
        f.write(f"LANG={config['locale_lang']}\n")
        f.write(f"LC_CTYPE={config['locale_ctype']}\n")
    try:
        subprocess.run(["chroot", "/target", "locale-gen"], check=True)
    except Exception as e:
        fallback(e)

def write_timezone(timezone=config.get("timezone", "UTC")):
    with open("/target/etc/timezone", "w") as f:
        f.write(f"{timezone}\n")
    if os.path.islink("/target/etc/localtime") or os.path.exists("/target/etc/localtime"):
        os.unlink("/target/etc/localtime")
    os.symlink(f"../usr/share/zoneinfo/{timezone}", "/target/etc/localtime")

def write_hostname(hostname=config.get("hostname", "localhost")):
    with open("/target/etc/hostname", "w") as f:
        f.write(f"{hostname}\n")
    with open("/target/etc/hosts", "w") as f:
        f.write("127.0.0.1 localhost\n")
        f.write(f"127.0.0.1 {hostname}\n")
        f.write("\n")
        f.write("# The following lines are desirable for IPv6 capable hosts\n")
        f.write("::1     localhost ip6-localhost ip6-loopback\n")
        f.write("ff02::1 ip6-allnodes\n")
        f.write("ff02::2 ip6-allrouters\n")

def write_sources_list():
    with open("/target/etc/apt/sources.list", "w") as f:
        f.write("\n".join(config.get("repo", [])) + "\n")

def run_hooks():
    try:
        for hook in config.get("hooks", []):
            subprocess.run(["chroot", "/target"] + hook, check=True)
    except Exception as e:
        fallback(e)

def create_user():
    try:
        for user in config.get("users", {}).keys():
            subprocess.run(
                ["chroot", "/target", "useradd",
                    "-m", user,
                    "-c", config["users"][user].get("realname", "Linux User"),
                    "-G", ",".join(config["users"][user].get("groups", [])),
                    "-s", "/bin/bash",
                    "-p", config["users"][user].get("passwd", "Zurna Dürüm"),
                    "-U"
                ]
            )
    except Exception as e:
        fallback(e)

def umount():
    dirs = ["/source"]
    with open("/proc/mounts", "r") as f:
        for line in f.read().strip().split("\n"):
            target = line.split(" ")[1]
            if "/target" in target:
                dirs.append(target)
    dirs.sort()
    dirs.reverse()
    for dir in dirs:
        subprocess.run(["umount", "-lf", dir], check=False)

def quit(status):
    subprocess.run(["sync"], check=False)
    if os.getpid() == 1:
        i = 5
        while i > 0:
            print(f"System will reboot in {i} seconds")
            time.sleep(1)
            i = i - 1
        with open("/proc/sys/kernel/sysrq", "w") as f:
            f.write("1")
        with open("/proc/sysrq-trigger", "w") as f:
            f.write("_sub")
    else:
        sys.exit(status)

if __name__ == "__main__":
    if os.getpid() == 1:
        cttyhack()
    # detect
    disk, prefix = detect_primary_disk()
    if not disk:
        fallback(RuntimeError("Could not find a suitable disk."))
    # modify disk & mount
    create_parts(disk, prefix)
    # copy filesystem
    rsync()
    # post copy
    bind()
    remove_live()
    install_grub(disk)
    # configure
    write_fstab(disk, prefix)
    write_x11_keyboard()
    write_language()
    write_timezone()
    write_hostname()
    write_sources_list()
    create_user()
    run_hooks()
    # final
    bind(unbind=True)
    umount()
    quit(0)
