#!/usr/bin/env python3
"""
Generate setup_bg.bin — the 800×480 4bpp background for the device setup screen.
The QR code is overlaid at runtime at position (QR_X, QR_Y) by the firmware.
Run from the firmware/ directory: python3 scripts/gen_setup_bg.py
"""

from PIL import Image, ImageDraw, ImageFont
import struct, os, sys

# ── Display + palette ───────────────────────────────────────────────────────────
W, H = 800, 480

# EPD 4bpp palette nibbles
BLACK  = 0x0
WHITE  = 0x1
YELLOW = 0x2
RED    = 0x3
BLUE   = 0x5
GREEN  = 0x6

# PIL RGB for each nibble (used for drawing and for quantisation)
PALETTE_RGB = {
    BLACK:  (0,   0,   0  ),
    WHITE:  (255, 255, 255),
    YELLOW: (255, 230, 0  ),
    RED:    (200, 0,   0  ),
    BLUE:   (0,   0,   220),
    GREEN:  (0,   170, 60 ),
}

# ── QR code placement (must match SETUP_QR_X / SETUP_QR_Y in epd.cpp) ──────────
QR_CELL = 5
QR_MODS = 41          # version 6, ECC_LOW
QR_PX   = QR_MODS * QR_CELL   # 205 px
QR_X    = 555
QR_Y    = (H - QR_PX) // 2    # 137

# ── Fonts ────────────────────────────────────────────────────────────────────────
FONT_DIR = "/usr/share/fonts/truetype/dejavu"
def font(name, size):
    try:
        return ImageFont.truetype(os.path.join(FONT_DIR, name), size)
    except Exception:
        return ImageFont.load_default()

font_title  = font("DejaVuSans-Bold.ttf", 36)
font_label  = font("DejaVuSans-Bold.ttf", 20)
font_sub    = font("DejaVuSans.ttf",      15)
font_scan   = font("DejaVuSans.ttf",      14)

# ── Draw ─────────────────────────────────────────────────────────────────────────
img  = Image.new("RGB", (W, H), PALETTE_RGB[WHITE])
draw = ImageDraw.Draw(img)

BK = PALETTE_RGB[BLACK]
GR = PALETTE_RGB[GREEN]

# Title
draw.text((40, 32), "pictureFrame", font=font_title, fill=BK)

# Thin rule under title
draw.rectangle([40, 80, 490, 82], fill=BK)

# ── Landscape diagram ────────────────────────────────────────────────────────────
LS_X, LS_Y, LS_W, LS_H = 45, 110, 200, 120
RIB_W, RIB_H = 56, 14
LS_RX = LS_X + (LS_W - RIB_W) // 2
LS_RY = LS_Y + LS_H   # ribbon protrudes below

BORDER = 3
draw.rectangle([LS_X, LS_Y, LS_X+LS_W, LS_Y+LS_H], outline=BK, width=BORDER)
draw.rectangle([LS_RX, LS_RY, LS_RX+RIB_W, LS_RY+RIB_H], fill=GR)

draw.text((LS_X, LS_RY + RIB_H + 10), "Landscape",     font=font_label, fill=BK)
draw.text((LS_X, LS_RY + RIB_H + 34), "Ribbon at bottom", font=font_sub, fill=BK)

# ── Portrait diagram ──────────────────────────────────────────────────────────────
PT_X, PT_Y, PT_W, PT_H = 310, 95, 120, 200
RIB2_W, RIB2_H = 14, 56
PT_RX = PT_X - RIB2_W   # ribbon protrudes left
PT_RY = PT_Y + (PT_H - RIB2_H) // 2

draw.rectangle([PT_X, PT_Y, PT_X+PT_W, PT_Y+PT_H], outline=BK, width=BORDER)
draw.rectangle([PT_RX, PT_RY, PT_RX+RIB2_W, PT_RY+RIB2_H], fill=GR)

draw.text((PT_X, PT_Y + PT_H + 10), "Portrait",       font=font_label, fill=BK)
draw.text((PT_X, PT_Y + PT_H + 34), "Ribbon on left", font=font_sub,   fill=BK)

# ── Divider ───────────────────────────────────────────────────────────────────────
draw.rectangle([520, 0, 522, H], fill=PALETTE_RGB[BLACK])

# ── QR zone label ─────────────────────────────────────────────────────────────────
scan_txt = "Scan to set up"
bb = draw.textbbox((0, 0), scan_txt, font=font_scan)
tw = bb[2] - bb[0]
draw.text((QR_X + (QR_PX - tw) // 2, QR_Y + QR_PX + 12), scan_txt, font=font_scan, fill=BK)

# Leave QR area pure WHITE so the firmware overlay is clean
draw.rectangle([QR_X, QR_Y, QR_X + QR_PX - 1, QR_Y + QR_PX - 1], fill=PALETTE_RGB[WHITE])

# ── Quantise to EPD palette ───────────────────────────────────────────────────────
def nearest(r, g, b):
    best, best_d = WHITE, float("inf")
    for nibble, (pr, pg, pb) in PALETTE_RGB.items():
        d = (r-pr)**2 + (g-pg)**2 + (b-pb)**2
        if d < best_d:
            best, best_d = nibble, d
    return best

pixels = img.load()
out = bytearray()
for y in range(H):
    for x in range(0, W, 2):
        hi = nearest(*pixels[x,   y])
        lo = nearest(*pixels[x+1, y])
        out.append((hi << 4) | lo)

out_path = os.path.join(os.path.dirname(__file__), "../data/setup_bg.bin")
with open(out_path, "wb") as f:
    f.write(out)

print(f"Written {len(out):,} bytes → {os.path.abspath(out_path)}")
print(f"QR overlay position: x={QR_X}, y={QR_Y}, cell={QR_CELL}  (copy to SETUP_QR_* in epd.cpp)")

# ── Preview PNG (for inspection) ─────────────────────────────────────────────────
preview = Image.new("RGB", (W, H))
pix = preview.load()
for y in range(H):
    for x in range(0, W, 2):
        byte = out[y * (W // 2) + x // 2]
        pix[x,   y] = PALETTE_RGB.get(byte >> 4,  (128,128,128))
        pix[x+1, y] = PALETTE_RGB.get(byte & 0xF, (128,128,128))
preview_path = out_path.replace(".bin", "_preview.png")
preview.save(preview_path)
print(f"Preview PNG → {os.path.abspath(preview_path)}")