"""
AIM bkg.tgz codec — reverse-engineered 2026-05-11.

Format:
    File on disk = zlib.compress(<18-byte header> + <W*H*2 bytes ARGB1555 pixels>)

Header (18 bytes, all little-endian):
    bytes  0- 1 : 0x0000           (always — magic/version?)
    bytes  2- 3 : 0x0002           (format flag — probably 2 = ARGB1555)
    bytes  4-11 : 0x00 * 8         (zero padding)
    bytes 12-13 : width  (uint16)  (e.g. 800)
    bytes 14-15 : height (uint16)  (e.g. 480)
    bytes 16-17 : 0x2110           (unknown — possibly stride or color config)

Pixels:
    16-bit ARGB1555: AAAAA RRRRR GGGGG BBBBB
        bit 15    = alpha (0=transparent, 1=opaque)
        bits 14-10 = R (5 bits, 0..31, scale << 3 for 8-bit)
        bits  9- 5 = G (5 bits)
        bits  4- 0 = B (5 bits)
    Row-major (row 0 = top, x increases left-to-right).
    Pure transparent black = 0x0000. Pure opaque black = 0x8000.

Hardware-verified path (planned for v0.3.7):
    .zconf2 XML at PATH_IMAGE_SFONDO_FW = "bkg.tgz" — firmware reads this,
    NOT the sibling bkg.png. Replace bkg.tgz to change the dash background.
"""

import struct
import zlib
from pathlib import Path
from PIL import Image


def decode(path):
    """Returns (width, height, PIL Image in RGBA mode) from a bkg.tgz file."""
    raw = Path(path).read_bytes()
    data = zlib.decompress(raw)
    if len(data) < 18:
        raise ValueError("Decompressed payload too small")
    width = struct.unpack_from("<H", data, 12)[0]
    height = struct.unpack_from("<H", data, 14)[0]
    expected = 18 + width * height * 2
    if len(data) != expected:
        raise ValueError(
            f"Pixel count mismatch: header says {width}x{height} -> "
            f"{expected} bytes, got {len(data)}"
        )
    im = Image.new("RGBA", (width, height))
    pixels = im.load()
    for y in range(height):
        for x in range(width):
            off = 18 + (y * width + x) * 2
            p = struct.unpack_from("<H", data, off)[0]
            a = 255 if (p >> 15) & 1 else 0
            r = ((p >> 10) & 0x1F) << 3
            g = ((p >> 5) & 0x1F) << 3
            b = (p & 0x1F) << 3
            pixels[x, y] = (r, g, b, a)
    return width, height, im


def encode(image, out_path):
    """Encode a PIL Image to AIM bkg.tgz format.

    The image is converted to RGBA and resampled to its native size.
    Alpha < 128 becomes ARGB1555 alpha-bit 0 (transparent),
    otherwise alpha-bit 1 (opaque).
    """
    im = image.convert("RGBA")
    width, height = im.size
    header = struct.pack(
        "<HHQHHH",
        0x0000,        # bytes 0-1
        0x0002,        # bytes 2-3 (format flag)
        0,             # bytes 4-11 (8 zero bytes)
        width,         # bytes 12-13
        height,        # bytes 14-15
        0x2110,        # bytes 16-17 (unknown, match stock value)
    )
    assert len(header) == 18, f"Header length {len(header)} != 18"
    pixels = bytearray(width * height * 2)
    px = im.load()
    for y in range(height):
        for x in range(width):
            r, g, b, a = px[x, y]
            r5 = (r >> 3) & 0x1F
            g5 = (g >> 3) & 0x1F
            b5 = (b >> 3) & 0x1F
            a1 = 1 if a >= 128 else 0
            p = (a1 << 15) | (r5 << 10) | (g5 << 5) | b5
            struct.pack_into("<H", pixels, (y * width + x) * 2, p)
    payload = header + bytes(pixels)
    compressed = zlib.compress(payload, level=9)
    Path(out_path).write_bytes(compressed)
    return out_path


if __name__ == "__main__":
    import sys
    if len(sys.argv) < 2:
        print("usage:")
        print("  bkg_tgz.py decode <bkg.tgz> <out.png>")
        print("  bkg_tgz.py encode <in.png>  <bkg.tgz>")
        print("  bkg_tgz.py info   <bkg.tgz>")
        sys.exit(0)
    cmd = sys.argv[1]
    if cmd == "decode":
        w, h, im = decode(sys.argv[2])
        im.save(sys.argv[3])
        print(f"decoded {w}x{h} -> {sys.argv[3]}")
    elif cmd == "encode":
        im = Image.open(sys.argv[2])
        encode(im, sys.argv[3])
        print(f"encoded {im.size} -> {sys.argv[3]}")
    elif cmd == "info":
        raw = Path(sys.argv[2]).read_bytes()
        data = zlib.decompress(raw)
        w = struct.unpack_from("<H", data, 12)[0]
        h = struct.unpack_from("<H", data, 14)[0]
        print(f"compressed: {len(raw)} bytes")
        print(f"decompressed: {len(data)} bytes")
        print(f"header (18B): {data[:18].hex()}")
        print(f"width x height: {w} x {h}")
        print(f"format: ARGB1555, 16-bit per pixel")
