#!/usr/bin/env python3

import subprocess
import argparse
import shutil
import sys
import os


class GrubConfig:

    def __init__(self, kernel, initrd):
        self.kernel = kernel
        self.initrd = initrd

    def __str__(self):
        if self.initrd:
            return """
set timeout=0
set default="0"
insmod efi_gop
menuentry "kernel" {
    multiboot2 %s
    module2 %s
}
""" % (self.kernel, self.initrd)
        else:
            return """
set timeout=0
set default="0"
insmod efi_gop
menuentry "kernel" {
    multiboot2 %s
}
""" % self.kernel

def find_executable(executables):
    for e in executables:
        path = shutil.which(e)
        if path:
            return path
    return None

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("image", help="OS image to build")
    parser.add_argument("-d", "--build-directory",
                        type=str, default="build/iso")
    parser.add_argument("-k", "--kernel", type=str,
                        default="kernel.elf", help="kernel image")
    parser.add_argument("-i", "--initrd", type=str,
                        default=None, help="initrd image")
    args = parser.parse_args()

    boot_dir = "boot"

    # Grub directory on build directory:
    build_grub_dir = os.path.join(
        args.build_directory, os.path.join(boot_dir, "grub"))

    # Grub configuration filename on build directory:
    build_grub_cfg = os.path.join(build_grub_dir, "grub.cfg")

    # Boot directory on build directory:
    build_boot_dir = os.path.join(args.build_directory, boot_dir)

    # Boot directory on bootable image filesystem:
    image_boot_dir = os.path.join("/", boot_dir)

    # Kernel filename on the bootable image filesystem:
    image_kernel = os.path.join(image_boot_dir, os.path.basename(args.kernel))

    # Initrd filename on the bootable image filesystem:
    if args.initrd:
        image_initrd = os.path.join(
            image_boot_dir, os.path.basename(args.initrd))
    else:
        image_initrd = None

    grub_cfg = GrubConfig(image_kernel, image_initrd)

    # Make directories on bootable image:
    os.makedirs(build_grub_dir, exist_ok=True)
    os.makedirs(build_boot_dir, exist_ok=True)

    # Copy files to the bootable image build directory:
    with open(build_grub_cfg, "w") as f:
        f.write(str(grub_cfg))
    shutil.copy(args.kernel, build_boot_dir)
    if args.initrd:
        shutil.copy(args.initrd, build_boot_dir)

    grub_mkrescue = find_executable(["grub2-mkrescue", "grub-mkrescue"])
    if grub_mkrescue is None:
        sys.exit("no grub-mkrescue found")

    # Build the bootable image:
    devnull = open(os.devnull, 'wb')
    ret = subprocess.call([grub_mkrescue, "-o", args.image,
                           args.build_directory], stdout=devnull, stderr=devnull)
    if ret != 0:
        sys.exit("%s failed" % (grub_mkrescue))
