#!/usr/bin/env python3

from subprocess import run
import argparse
import os
import shutil
import hashlib

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
            formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('-k', '--kernel', required=True, help='path to kernel binary')
    parser.add_argument('-c', '--config', required=True, help='path to bootloader config')
    parser.add_argument('-o', '--output', required=True, help='path for output image')

    parser.add_argument('--prefix-dir', default='isodir', help='prefix dir')
    parser.add_argument('--boot-dir', default='boot/grub', help='bootloader dir')
    parser.add_argument('--kernel-dir', default='boot/glitch', help='kernel dir')

    # grub tools
    parser.add_argument('--grub-file', default='grub-file',
                        help='grub-file executable')
    parser.add_argument('--grub-script-check', default='grub-script-check',
                        help='grub-script-check executable')
    parser.add_argument('--grub-mkrescue', default='grub-mkrescue',
                        help='grub-mkrescue executable')
    # other tools
    parser.add_argument('--strip', default='i686-elf-strip',
                        help='strip executable')

    args = parser.parse_args()
    args.boot_dir = f'{args.prefix_dir}/{args.boot_dir}'
    args.kernel_dir = f'{args.prefix_dir}/{args.kernel_dir}'

    # check inputs
    run([ args.grub_file, '--is-x86-multiboot2', args.kernel ], check=True)
    run([ args.grub_script_check, args.config ], check=True)

    # create directory structure and copy files
    os.makedirs(args.boot_dir, exist_ok=True)
    os.makedirs(args.kernel_dir, exist_ok=True)
    shutil.copyfile(args.config, f'{args.boot_dir}/grub.cfg')
    shutil.copyfile(args.kernel, f'{args.kernel_dir}/glitch.elf')

    # strip and hash binary
    run([ args.strip, f'{args.kernel_dir}/glitch.elf' ], check=True)
    with open(f'{args.kernel_dir}/checksums', 'w') as checksums:
        with open(f'{args.kernel_dir}/glitch.elf', 'rb') as f:
            digest = hashlib.file_digest(f, 'sha512').hexdigest()
            path = f'{args.kernel_dir}/{args.kernel}'.removeprefix(args.prefix_dir)
            print(f'{digest}  {path}', file=checksums)

    # create iso image
    run([ args.grub_mkrescue, '-o', args.output, args.prefix_dir ], check=True)

    # cleanup
    #shutil.rmtree(args.prefix_dir)