#!/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)")
    return raw.ljust(FIELD, b"\x00")


def main() -> None:
    p = argparse.ArgumentParser(description="append the ninlight OTA trailer")
    p.add_argument("--image", required=True, type=Path, help="input _gz.img")
    p.add_argument("--device", required=True, help="device id, e.g. e27_rgb")
    p.add_argument("--version", required=True, help="firmware version")
    p.add_argument("--out", required=True, type=Path, help="output _ota.img")
    a = p.parse_args()

    data = a.image.read_bytes()
    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()