#!/usr/bin/env python3

import argparse
import binascii
from pathlib import Path


DEFAULT_FPGA_BITSTREAM = "FPGA/VNA/top.bin"
DEFAULT_OUTPUT = "combined.vnafw"
DEFAULT_MCU_CANDIDATES = [
    "Software/VNA_embedded/Debug/VNA_embedded.bin",
    "Software/VNA_embedded/Release/VNA_embedded.bin",
    "Software/VNA_embedded/build/VNA_embedded.bin",
]

HEADER_SIZE = 24


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Assemble a combined LibreVNA firmware image.",
    )
    parser.add_argument(
        "--bitstream",
        default=DEFAULT_FPGA_BITSTREAM,
        help=f"FPGA bitstream to package (default: {DEFAULT_FPGA_BITSTREAM})",
    )
    parser.add_argument(
        "--output",
        default=DEFAULT_OUTPUT,
        help=f"Output firmware image path (default: {DEFAULT_OUTPUT})",
    )
    parser.add_argument(
        "--mcu",
        help="Explicit MCU firmware .bin path. If omitted, the newest known build is used.",
    )
    return parser.parse_args()


def select_mcu_firmware(explicit_path: str | None) -> Path:
    if explicit_path is not None:
        path = Path(explicit_path)
        if not path.is_file():
            raise FileNotFoundError(f"MCU firmware not found: {path}")
        return path

    latest_modification = 0.0
    newest_mcu: Path | None = None
    for candidate in DEFAULT_MCU_CANDIDATES:
        path = Path(candidate)
        if not path.is_file():
            continue
        modified_time = path.stat().st_mtime
        if modified_time >= latest_modification:
            latest_modification = modified_time
            newest_mcu = path

    if newest_mcu is None:
        raise FileNotFoundError("Couldn't find MCU firmware file")
    return newest_mcu


def crc32_from_file(filename: Path, initial_crc: int) -> int:
    buffer = filename.read_bytes()
    return binascii.crc32(buffer, initial_crc) & 0xFFFFFFFF


def main() -> int:
    args = parse_args()

    bitstream_path = Path(args.bitstream)
    if not bitstream_path.is_file():
        raise FileNotFoundError(f"FPGA bitstream not found: {bitstream_path}")

    mcu_firmware_path = select_mcu_firmware(args.mcu)
    output_path = Path(args.output)

    print(f"Using {mcu_firmware_path} as MCU firmware")

    size_fpga = bitstream_path.stat().st_size
    size_mcu = mcu_firmware_path.stat().st_size
    print(f"Got FPGA bitstream of size {size_fpga}")
    print(f"Got MCU firmware of size {size_mcu}")

    crc = crc32_from_file(bitstream_path, 0xFFFFFFFF)
    crc = crc32_from_file(mcu_firmware_path, crc)
    print(f"Calculating CRC...:{hex(crc)}")

    output_path.parent.mkdir(parents=True, exist_ok=True)
    with output_path.open("wb") as firmware_image:
        firmware_image.write(b"VNA!")
        firmware_image.write((HEADER_SIZE).to_bytes(4, byteorder="little"))
        firmware_image.write(size_fpga.to_bytes(4, byteorder="little"))
        firmware_image.write((HEADER_SIZE + size_fpga).to_bytes(4, byteorder="little"))
        firmware_image.write(size_mcu.to_bytes(4, byteorder="little"))
        firmware_image.write(crc.to_bytes(4, byteorder="little"))

        if firmware_image.tell() != HEADER_SIZE:
            raise RuntimeError(
                "Incorrect header size "
                f"(defined as {HEADER_SIZE} but actual header is {firmware_image.tell()})"
            )

        firmware_image.write(bitstream_path.read_bytes())
        firmware_image.write(mcu_firmware_path.read_bytes())

        if firmware_image.tell() % 256 != 0:
            padding = 256 - firmware_image.tell() % 256
            firmware_image.write(bytearray(padding))

        final_size = firmware_image.tell()

    print(f"Created combined firmware file at {output_path} with size {final_size}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())