#!/usr/bin/env python3
"""Build the one-time WildFire image from server-built firmware artifacts."""

from __future__ import annotations

import argparse
import struct
import zlib
from pathlib import Path


BOOT_BASE = 0x10000000
APP_BASE = 0x10010000
APP_END = 0x10080000
METADATA_ADDRESS = 0x14000000
METADATA_MAGIC = 0x4F544131
METADATA_VERSION = 1


def parse_ihex(path: Path) -> tuple[dict[int, int], int | None]:
    memory: dict[int, int] = {}
    upper = 0
    start_linear: int | None = None
    for line_number, raw in enumerate(path.read_text(encoding="ascii").splitlines(), 1):
        line = raw.strip()
        if not line:
            continue
        if not line.startswith(":"):
            raise ValueError(f"{path}:{line_number}: invalid Intel HEX record")
        record = bytes.fromhex(line[1:])
        if len(record) < 5 or sum(record) & 0xFF:
            raise ValueError(f"{path}:{line_number}: invalid record checksum")
        length = record[0]
        if len(record) != length + 5:
            raise ValueError(f"{path}:{line_number}: invalid record length")
        offset = int.from_bytes(record[1:3], "big")
        kind = record[3]
        payload = record[4 : 4 + length]
        if kind == 0:
            address = upper + offset
            for index, value in enumerate(payload):
                absolute = address + index
                previous = memory.get(absolute)
                if previous is not None and previous != value:
                    raise ValueError(f"overlapping data at 0x{absolute:08X}")
                memory[absolute] = value
        elif kind == 2:
            if length != 2:
                raise ValueError("invalid extended segment record")
            upper = int.from_bytes(payload, "big") << 4
        elif kind == 4:
            if length != 2:
                raise ValueError("invalid extended linear record")
            upper = int.from_bytes(payload, "big") << 16
        elif kind == 5:
            if length != 4:
                raise ValueError("invalid start linear address record")
            start_linear = int.from_bytes(payload, "big")
        elif kind in (1, 3):
            continue
        else:
            raise ValueError(f"unsupported Intel HEX record type {kind}")
    if not memory:
        raise ValueError("bootloader HEX contains no data")
    return memory, start_linear


def record(address: int, kind: int, payload: bytes) -> str:
    body = bytes((len(payload), (address >> 8) & 0xFF, address & 0xFF, kind)) + payload
    checksum = (-sum(body)) & 0xFF
    return ":" + (body + bytes((checksum,))).hex().upper()


def write_ihex(path: Path, memory: dict[int, int], start_linear: int | None) -> None:
    lines: list[str] = []
    current_upper: int | None = None
    addresses = sorted(memory)
    index = 0
    while index < len(addresses):
        start = addresses[index]
        upper = start >> 16
        if upper != current_upper:
            lines.append(record(0, 4, upper.to_bytes(2, "big")))
            current_upper = upper
        low = start & 0xFFFF
        payload = bytearray((memory[start],))
        index += 1
        while index < len(addresses) and len(payload) < 16:
            next_address = addresses[index]
            if next_address != start + len(payload) or next_address >> 16 != upper:
                break
            payload.append(memory[next_address])
            index += 1
        lines.append(record(low, 0, bytes(payload)))
    if start_linear is not None:
        lines.append(record(0, 5, start_linear.to_bytes(4, "big")))
    lines.append(record(0, 1, b""))
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text("\n".join(lines) + "\n", encoding="ascii", newline="\n")


def add_region(memory: dict[int, int], address: int, data: bytes, label: str) -> None:
    for index, value in enumerate(data):
        absolute = address + index
        if absolute in memory:
            raise ValueError(f"{label} overlaps existing data at 0x{absolute:08X}")
        memory[absolute] = value


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--bootloader", type=Path, required=True)
    parser.add_argument("--calibration", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()

    memory, start_linear = parse_ihex(args.bootloader)
    if min(memory) < BOOT_BASE or max(memory) >= APP_BASE:
        raise ValueError("bootloader data is outside the immutable 64 KiB region")

    calibration = args.calibration.read_bytes()
    if len(calibration) < 8 or APP_BASE + len(calibration) > APP_END:
        raise ValueError("calibration firmware size is outside the application region")
    crc32 = zlib.crc32(calibration) & 0xFFFFFFFF
    metadata = struct.pack(
        "<8I",
        METADATA_MAGIC,
        METADATA_VERSION,
        len(calibration),
        crc32,
        (~len(calibration)) & 0xFFFFFFFF,
        (~crc32) & 0xFFFFFFFF,
        APP_BASE,
        (~APP_BASE) & 0xFFFFFFFF,
    )
    add_region(memory, APP_BASE, calibration, "calibration firmware")
    add_region(memory, METADATA_ADDRESS, metadata, "boot metadata")
    write_ihex(args.output, memory, start_linear)
    print(
        "FACTORY_IMAGE_OK "
        f"boot_bytes={sum(BOOT_BASE <= value < APP_BASE for value in memory)} "
        f"app_bytes={len(calibration)} crc32=0x{crc32:08X} output={args.output}"
    )
    return 0


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