#!/usr/bin/python3 -I
import os
import subprocess
import sys
import re
from typing import List, Tuple

def main(args):
    section_re = re.compile(rb"\A *(?:0|[1-9][0-9]*) +([!-~]+) +([0-9a-f]{8}) +([0-9a-f]{16}) +[0-9a-f]{16} +[0-9a-f]{8} +2")
    alignment_mask = (1 << 21) - 1
    if len(args) != 6:
        print(f"Usage: uki-generate HYPERVISOR CONFIG KERNEL INITRAMFS OUTPUT", file=sys.stderr)
        sys.exit(1)
    _, hyp, cfg, kern, initramfs, out = args
    if hyp[0] != '/':
        hyp = './' + hyp
    if out[0] != '/':
        out = './' + out
    output = subprocess.check_output([
        "objdump",
        "-WE", # do not use debuginfod
        "--section-headers",
        "--",
        hyp,
    ])
    max_vma = 0
    for line in output.splitlines():
        m = section_re.match(line)
        if not m:
            continue
        section_name, size, start_vma = m.group(1), int(m.group(2), 16), int(m.group(3), 16)
        if section_name.startswith(b".annobin"):
            continue
        max_vma = max(max_vma, size, + start_vma)
    def round_to_next(f: int) -> int:
        max_address = (0xffffffffffffffff & ~alignment_mask)
        if f > max_address:
            print(f"Fatal error: Address overflow: {f} exceeds {max_address}", file=sys.stderr)
            sys.exit(1)
        return (f + alignment_mask) & ~alignment_mask
    base_address = round_to_next(max_vma)
    kernel_vma = round_to_next(base_address + os.stat(cfg).st_size)
    initramfs_vma = round_to_next(kernel_vma + os.stat(kern).st_size)
    cmdline = [
        "objcopy",
        f"--section-alignment={alignment_mask + 1}",
        f"--file-alignment={1 << 5}",
        #"--remove-section=.buildid",
        "--remove-section=.annobin.*",
        #"--strip-debug",
        f"--add-section=.config={cfg}",
        f"--change-section-vma=.config={base_address}",
        f"--add-section=.kernel={kern}",
        f"--change-section-vma=.kernel={kernel_vma}",
        f"--add-section=.ramdisk={initramfs}",
        f"--change-section-vma=.ramdisk={initramfs_vma}",
        "--long-section-names=disable",
        "--",
        hyp,
        out,
    ]
    subprocess.check_call(cmdline, stdin=subprocess.DEVNULL, stdout=subprocess.DEVNULL)
if __name__ == '__main__':
    try:
        main(sys.argv)
    except subprocess.CalledProcessError as e:
        sys.exit(e.returncode)
