#!/usr/bin/env python3
"""Appends the ninlight OTA trailer to a `_gz.img` produced by wm_tool.

Layout (512 bytes): "NINL" | version(1) | device[16] | fw_version[16] | sha256[32] | 0xFF padding.
The hash covers the complete `_gz.img` (T_BOOTER header + payload), which is exactly what the
firmware hashes while writing the image to flash.
"""
import argparse
import hashlib
import sys
from pathlib import Path

TRAILER_SIZE = 512
TRAILER_VERSION = 1
FIELD = 16


def field(value: str, name: str) -> bytes:
    raw = value.encode()
    if len(raw) >= FIELD:
        sys.exit(f"error: {name} '{value}' is too long (max {FIELD - 1} bytes)")
    if any(c.isspace() for c in value):
        sys.exit(f"error: {name} '{value}' contains whitespace")
    return raw.ljust(FIELD, b"\x00")


def verify(path: Path, expect_device: str | None) -> int:
    """Checks an _ota.img the way the firmware does before accepting it."""
    try:
        data = path.read_bytes()
    except OSError as e:
        print(f"error: {e}", file=sys.stderr)
        return 1
    if len(data) <= TRAILER_SIZE:
        print(f"error: {path} is too short to carry a trailer", file=sys.stderr)
        return 1
    payload, trailer = data[:-TRAILER_SIZE], data[-TRAILER_SIZE:]
    if trailer[0:4] != b"NINL":
        print(f"error: {path} has no ninlight trailer", file=sys.stderr)
        return 1
    if trailer[4] != TRAILER_VERSION:
        print(f"error: {path} has trailer version {trailer[4]}, expected {TRAILER_VERSION}",
              file=sys.stderr)
        return 1
    # Device field must be null-terminated within its 16 bytes (same as firmware check)
    device_field = trailer[5:21]
    if b"\x00" not in device_field:
        print(f"error: {path} has unterminated device field", file=sys.stderr)
        return 1
    device = device_field.rstrip(b"\x00").decode("ascii", "replace")
    version = trailer[21:37].rstrip(b"\x00").decode("ascii", "replace")
    if expect_device is not None and device != expect_device:
        print(f"error: {path} carries device '{device}', expected '{expect_device}'",
              file=sys.stderr)
        return 1
    if trailer[37:69] != hashlib.sha256(payload).digest():
        print(f"error: {path} fails its sha256 - the payload does not match the trailer",
              file=sys.stderr)
        return 1
    print(f"{device} {version} {len(data)}")
    return 0


def main() -> None:
    p = argparse.ArgumentParser(description="append the ninlight OTA trailer")
    p.add_argument("--image", type=Path, help="input _gz.img")
    p.add_argument("--device", help="device id, e.g. e27_rgb")
    p.add_argument("--version", help="firmware version")
    p.add_argument("--out", type=Path, help="output _ota.img")
    p.add_argument("--verify", type=Path, help="check an existing _ota.img instead of building one")
    p.add_argument("--expect-device", help="with --verify: require this device id")
    a = p.parse_args()

    if a.verify is not None:
        # Verify mode: check that build arguments are not provided
        build_args = ("image", "device", "version", "out")
        if any(getattr(a, name) is not None for name in build_args):
            p.error("--verify cannot be combined with build mode arguments (--image, --device, --version, --out)")
        raise SystemExit(verify(a.verify, a.expect_device))

    # Build mode: require all arguments
    for name in ("image", "device", "version", "out"):
        if getattr(a, name) is None:
            p.error(f"--{name} is required unless --verify is given")

    try:
        data = a.image.read_bytes()
    except OSError as e:
        sys.exit(f"error: {e}")

    if len(data) > TRAILER_SIZE and data[-TRAILER_SIZE : -TRAILER_SIZE + 4] == b"NINL":
        sys.exit("error: image already carries a trailer")

    trailer = bytearray(b"\xff" * TRAILER_SIZE)
    trailer[0:4] = b"NINL"
    trailer[4] = TRAILER_VERSION
    trailer[5:21] = field(a.device, "device")
    trailer[21:37] = field(a.version, "version")
    trailer[37:69] = hashlib.sha256(data).digest()
    a.out.write_bytes(data + bytes(trailer))
    print(f"{a.out} ({len(data) + TRAILER_SIZE} bytes, device={a.device}, version={a.version})")


if __name__ == "__main__":
    main()