#!/usr/bin/env python3
import argparse
import json
import re
import struct
import subprocess
import zlib
from pathlib import Path


BOOT_IMAGE_MAGIC = 0x54425652
BOOT_IMAGE_VERSION = 1
BOOT_HEADER_SIZE = 64


def parse_int(value):
    return int(value, 0)


def run(cmd):
    return subprocess.check_output(cmd, text=True)


def elf_entry(readelf, elf):
    text = run([readelf, "-h", str(elf)])
    match = re.search(r"Entry point address:\s+(0x[0-9a-fA-F]+)", text)
    if not match:
        raise SystemExit(f"could not parse ELF entry from {elf}")
    return int(match.group(1), 16)


def elf_load_addr(readelf, elf):
    text = run([readelf, "-l", str(elf)])
    pattern = re.compile(
        r"LOAD\s+0x[0-9a-fA-F]+\s+0x[0-9a-fA-F]+\s+"
        r"(0x[0-9a-fA-F]+)\s*\n\s+"
        r"(0x[0-9a-fA-F]+)\s+(0x[0-9a-fA-F]+)",
        re.MULTILINE,
    )
    matches = list(pattern.finditer(text))
    if not matches:
        raise SystemExit(f"could not parse LOAD segment from {elf}")
    return min(int(match.group(1), 16) for match in matches)


def main():
    parser = argparse.ArgumentParser(description="Create VCU118 Rocket flash boot image")
    parser.add_argument("--elf", required=True)
    parser.add_argument("--out", required=True)
    parser.add_argument("--payload-out")
    parser.add_argument("--manifest-out")
    parser.add_argument("--flash-offset", type=parse_int, required=True)
    parser.add_argument("--load-addr", type=parse_int)
    parser.add_argument("--entry", type=parse_int)
    parser.add_argument("--objcopy", default="/tools/toolchains/rvtc/bin/riscv64-unknown-elf-objcopy")
    parser.add_argument("--readelf", default="/tools/toolchains/rvtc/bin/riscv64-unknown-elf-readelf")
    args = parser.parse_args()

    elf = Path(args.elf)
    out = Path(args.out)
    payload_out = Path(args.payload_out) if args.payload_out else out.with_suffix(".payload.bin")
    manifest_out = Path(args.manifest_out) if args.manifest_out else out.with_suffix(".json")

    out.parent.mkdir(parents=True, exist_ok=True)
    payload_out.parent.mkdir(parents=True, exist_ok=True)
    manifest_out.parent.mkdir(parents=True, exist_ok=True)

    subprocess.check_call([args.objcopy, "-O", "binary", str(elf), str(payload_out)])
    payload = payload_out.read_bytes()

    entry = args.entry if args.entry is not None else elf_entry(args.readelf, elf)
    load_addr = args.load_addr if args.load_addr is not None else elf_load_addr(args.readelf, elf)
    payload_flash_offset = args.flash_offset + BOOT_HEADER_SIZE
    crc = zlib.crc32(payload) & 0xFFFFFFFF

    header_words = [
        BOOT_IMAGE_MAGIC,
        BOOT_IMAGE_VERSION,
        BOOT_HEADER_SIZE,
        0,
        payload_flash_offset,
        load_addr,
        entry,
        len(payload),
        crc,
        0, 0, 0, 0, 0, 0, 0,
    ]
    out.write_bytes(struct.pack("<16I", *header_words) + payload)

    manifest = {
        "elf": str(elf),
        "image": str(out),
        "payload": str(payload_out),
        "flash_offset": f"0x{args.flash_offset:08x}",
        "payload_flash_offset": f"0x{payload_flash_offset:08x}",
        "load_addr": f"0x{load_addr:08x}",
        "entry": f"0x{entry:08x}",
        "size": len(payload),
        "crc32": f"0x{crc:08x}",
    }
    manifest_out.write_text(json.dumps(manifest, indent=2) + "\n")

    print(f"image: {out}")
    print(f"flash_offset: 0x{args.flash_offset:08x}")
    print(f"payload_flash_offset: 0x{payload_flash_offset:08x}")
    print(f"load_addr: 0x{load_addr:08x}")
    print(f"entry: 0x{entry:08x}")
    print(f"size: {len(payload)}")
    print(f"crc32: 0x{crc:08x}")


if __name__ == "__main__":
    main()