#!/usr/bin/env python3
"""
Generate ap_bg.bin and setup_bg.bin — 800×480 4bpp backgrounds for the
pictureFrame e-ink device.  QR overlay areas are left WHITE so the
firmware can render the actual QR code at runtime.

Run from the firmware/ directory:
    python3 scripts/gen_screens.py

Constants exported (copy to epd.cpp):
    AP_QR_X, AP_QR_Y, AP_QR_CELL, AP_QR_PX
    SETUP_QR_X, SETUP_QR_Y, SETUP_QR_CELL, SETUP_QR_PX
"""

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

# URL for the user manual QR baked into Step 1/2. Static on purpose —
# changing it requires regenerating the bg images + uploadfs.
MANUAL_URL = "https://pictureframe.edholm.me/help"

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

# ── EPD palette ───────────────────────────────────────────────────────────────
BLACK  = 0x0;  BK = (26,  26,  26 )
WHITE  = 0x1;  WH = (245, 245, 240)
YELLOW = 0x2;  YL = (240, 208, 0  )
RED    = 0x3;  RD = (192, 48,  32 )
BLUE   = 0x5;  BL = (24,  64,  192)
GREEN  = 0x6;  GR = (16,  160, 64 )

PALETTE_RGB = {BLACK: BK, WHITE: WH, YELLOW: YL, RED: RD, BLUE: BL, GREEN: GR}

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

def pack(img):
    """Convert RGB PIL image → 4bpp packed bytearray."""
    px = img.load()
    out = bytearray()
    for y in range(H):
        for x in range(0, W, 2):
            hi = nearest(*px[x,   y])
            lo = nearest(*px[x+1, y])
            out.append((hi << 4) | lo)
    return out

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

F_HEAD  = ttf("DejaVuSans-Bold.ttf", 26)
F_BAR   = ttf("DejaVuSans-Bold.ttf", 13)
F_STEP  = ttf("DejaVuSans.ttf",      13)
F_STEP_B= ttf("DejaVuSans-Bold.ttf", 13)
F_STEPN = ttf("DejaVuSans-Bold.ttf", 13)
F_LABEL = ttf("DejaVuSans-Bold.ttf", 11)
F_TINY  = ttf("DejaVuSans-Bold.ttf", 10)
F_FOOT  = ttf("DejaVuSans.ttf",      12)
F_CHIP  = ttf("DejaVuSans-Bold.ttf", 12)
F_SUB   = ttf("DejaVuSans.ttf",      14)
F_BIG   = ttf("DejaVuSans-Bold.ttf", 14)

# ── Layout constants ──────────────────────────────────────────────────────────
BAR_H    = 52
BODY_Y   = BAR_H          # 52

LEFT_X   = 0;   LEFT_W  = 310
DIV1_X   = 310; DIV_W   = 2
CTR_X    = 312; CTR_W   = 196
DIV2_X   = 508
RIGHT_X  = 510; RIGHT_W = 290   # 800-510

# QR positions (MUST match epd.cpp constants)
# WiFi-join QR — drawn at runtime by firmware. Cell shrunk from 5 to 4
# (148 px instead of 185 px) to leave room for a second QR below it that
# opens Safari → forces iOS captive UI.
AP_QR_CELL  = 4
AP_QR_MODS  = 37          # version 5, ECC_LOW
AP_QR_PX    = AP_QR_MODS * AP_QR_CELL   # 148

SETUP_QR_CELL = 5
SETUP_QR_MODS = 41        # version 6, ECC_LOW
SETUP_QR_PX   = SETUP_QR_MODS * SETUP_QR_CELL  # 205

# Centre of right panel
RIGHT_CX = RIGHT_X + RIGHT_W // 2   # 655
BODY_CY  = BODY_Y + (H - BODY_Y) // 2  # 266

# Stacked layout: WiFi QR (top) + URL QR (bottom). Each has a label
# above it. AP_QR is dynamic — firmware overlays it at runtime. URL QR
# is static (always http://192.168.4.1/) and baked into the bg image.
AP_QR_X     = RIGHT_CX - AP_QR_PX // 2   # centered horizontally
AP_QR_Y     = 100                          # below STEP 1 label
URL_QR_BOX  = 4                            # cell size in px
URL_QR_TARGET_Y = 320                      # below STEP 2 label

SETUP_QR_X = RIGHT_CX - SETUP_QR_PX // 2  # 553
SETUP_QR_Y = BODY_CY  - SETUP_QR_PX // 2  # 164
SETUP_QR_Y = 175    # nudge for label

def leave_qr_white(draw, qr_x, qr_y, qr_px):
    """Blank the QR overlay region so firmware can write the real QR."""
    draw.rectangle([qr_x, qr_y, qr_x+qr_px-1, qr_y+qr_px-1], fill=WH)

def text_center(draw, cx, y, text, font, fill):
    bb = draw.textbbox((0,0), text, font=font)
    tw = bb[2]-bb[0]
    draw.text((cx - tw//2, y), text, font=font, fill=fill)

def up_arrow(draw, cx, cy, half_w=12, h=22, color=BK):
    """Solid filled triangle pointing up, centered on (cx, cy)."""
    draw.polygon([
        (cx,           cy - h // 2),  # tip
        (cx - half_w,  cy + h // 2),  # base left
        (cx + half_w,  cy + h // 2),  # base right
    ], fill=color)


def left_arrow(draw, cx, cy, half_h=12, w=22, color=BK):
    """Solid filled triangle pointing left, centered on (cx, cy)."""
    draw.polygon([
        (cx - w // 2,  cy),            # tip
        (cx + w // 2,  cy - half_h),   # base top
        (cx + w // 2,  cy + half_h),   # base bottom
    ], fill=color)


def orientation_diagrams(draw, accent, show_active_ls=True):
    """Draw both orientation diagrams in the centre panel.

    The accent / show_active_ls parameters are kept for call-site
    compatibility but no longer drive any colour decisions — both
    diagrams render in black so neither orientation looks privileged.
    The supported orientation is communicated by the up-arrow inside
    the landscape screen instead."""
    cx = CTR_X + CTR_W // 2  # 410

    # ── Section title ─────────────────────────────────────────────
    text_center(draw, cx, BODY_Y+15, "FRAME",       F_TINY, BK)
    text_center(draw, cx, BODY_Y+27, "ORIENTATION", F_TINY, BK)

    # ── Landscape ──────────────────────────────────────────────────
    ls_x, ls_y, ls_w, ls_h = CTR_X+43, BODY_Y+52, 110, 66
    rib_w, rib_h = 110, 10

    text_center(draw, cx, ls_y-14, "LANDSCAPE", F_LABEL, BK)

    draw.rectangle([ls_x, ls_y, ls_x+ls_w-1, ls_y+ls_h-1], outline=BK, width=3)
    draw.rectangle([ls_x, ls_y+ls_h, ls_x+rib_w-1, ls_y+ls_h+rib_h-1], fill=BK)

    # Up arrow inside both diagrams — shows the user which edge is "up"
    # when they hang the frame, regardless of orientation.
    up_arrow(draw, ls_x + ls_w // 2, ls_y + ls_h // 2, half_w=14, h=26)

    # Thin separator
    sep_y = ls_y + ls_h + rib_h + 14
    draw.rectangle([CTR_X+10, sep_y, CTR_X+CTR_W-10, sep_y], fill=(180,180,175))

    # ── Portrait — drawn 90° CCW from upright so it rotates back to a
    # correct portrait view (tall rect, ribbon on left, arrow pointing
    # up) when the user tilts the frame 90° CW from landscape. In the
    # source/landscape view it appears as a wide rect with the ribbon
    # on the bottom and the arrow pointing left — looks "rotated"
    # because it is, on purpose.
    pt_w, pt_h = 106, 64
    pr_w, pr_h = 106, 10
    pt_x = CTR_X + (CTR_W - pt_w) // 2
    pt_y = sep_y + 14

    text_center(draw, cx, pt_y-14, "PORTRAIT", F_LABEL, BK)

    draw.rectangle([pt_x, pt_y, pt_x+pt_w-1, pt_y+pt_h-1], outline=BK, width=3)
    draw.rectangle([pt_x, pt_y+pt_h, pt_x+pr_w-1, pt_y+pt_h+pr_h-1], fill=BK)

    left_arrow(draw, pt_x + pt_w // 2, pt_y + pt_h // 2, half_h=12, w=22)


# ═══════════════════════════════════════════════════════════════════════════════
# AP SCREEN — accent-colored, WiFi credentials
# Pass accent=YL/header="SETUP MODE — STEP 1 OF 2"/qr_label="SCAN TO CONNECT"
# for the normal first-attempt screen, or accent=RD/header="CONNECTION FAILED
# — TRY AGAIN"/qr_label="Connection Failed — try again" for the post-WiFi-fail
# retry screen. Same layout either way so the panel diff is just color +
# header/label text.
# ═══════════════════════════════════════════════════════════════════════════════
def gen_ap(accent=YL, header="SETUP MODE  —  STEP 1 OF 2", qr_label="SCAN TO CONNECT"):
    img  = Image.new("RGB", (W, H), WH)
    draw = ImageDraw.Draw(img)

    # ── Status bar ────────────────────────────────────────────────
    draw.rectangle([0, 0, W-1, BAR_H-1], fill=accent)
    draw.text((24, 18), header, font=F_BAR, fill=BK)

    # Right chip: black box with device SSID
    chip_x, chip_y = 498, 11
    chip_text = "PictureFrame-91F8"
    bb = draw.textbbox((0,0), chip_text, font=F_CHIP)
    chip_w = bb[2]-bb[0] + 22
    chip_x2 = chip_x + chip_w
    draw.rectangle([chip_x, chip_y, chip_x2, BAR_H-12], fill=BK)
    draw.text((chip_x+11, chip_y+7), chip_text, font=F_CHIP, fill=accent)

    # ── Panel dividers ────────────────────────────────────────────
    draw.rectangle([DIV1_X, BODY_Y, DIV1_X+1, H-1], fill=BK)
    draw.rectangle([DIV2_X, BODY_Y, DIV2_X+1, H-1], fill=BK)

    # ── Left panel ────────────────────────────────────────────────
    # Heading
    draw.text((28, BODY_Y+20), "Connect to", font=F_HEAD, fill=BK)
    draw.text((28, BODY_Y+52), "WiFi",        font=F_HEAD, fill=BK)
    bb = draw.textbbox((0,0), "WiFi", font=F_HEAD)
    draw.rectangle([28, BODY_Y+82, 28+bb[2]+2, BODY_Y+85], fill=accent)

    # Steps — frames ship with battery unplugged to preserve shelf life
    # (idle setup-screen polling is non-trivial draw on e-ink), so the
    # very first prompt is "Plug in power". Step 2 is unlock-first (iOS
    # won't fire the captive UI from a locked-phone scan). Two-QR flow
    # because iOS in recent versions doesn't auto-open the captive
    # portal even after CNA detects it; scanning the second QR opens
    # Safari which forces the portal to render.
    steps = [
        ("Plug in power",              ""),
        ("Unlock your phone first",    ""),
        ("Scan QR 1 →",                "joins PictureFrame WiFi"),
        ("Scan QR 2 →",                "page opens in Safari"),
        ("Enter your WiFi password",   "and tap Connect"),
    ]
    sy = BODY_Y + 95
    step_pitch = 32
    for i, (l1, l2) in enumerate(steps):
        bx, by = 28, sy + i*step_pitch
        draw.rectangle([bx, by, bx+24, by+24], fill=BK)
        text_center(draw, bx+12, by+6, str(i+1), F_STEPN, accent)
        draw.text((62, by+3),  l1, font=F_STEP,   fill=BK)
        if l2:
            draw.text((62, by+17), l2, font=F_STEP, fill=BK)

    # Manual QR + side label — bottom of left panel.
    # The manual covers the captive-portal-didn't-open fallback (192.168.4.1)
    # plus general setup/troubleshooting, so we drop the inline footnote
    # and rely on the QR/manual to deliver the longer story.
    qr_bottom_y = BODY_Y + 95 + len(steps)*step_pitch + 12  # below last step
    qr = qrcode.QRCode(
        version=None,
        error_correction=qrcode.constants.ERROR_CORRECT_L,
        box_size=4,
        border=2,
    )
    qr.add_data(MANUAL_URL)
    qr.make(fit=True)
    qr_img = qr.make_image(fill_color=BK, back_color=WH).convert("RGB")
    qr_w, qr_h = qr_img.size
    img.paste(qr_img, (28, qr_bottom_y))

    # Side label, vertically centered against the QR
    label_x = 28 + qr_w + 14
    label_y = qr_bottom_y + (qr_h - 56) // 2
    draw.text((label_x, label_y),      "Need help?",        font=F_STEP_B, fill=BK)
    draw.text((label_x, label_y + 18), "Scan for setup",    font=F_STEP,   fill=BK)
    draw.text((label_x, label_y + 32), "& troubleshooting", font=F_STEP,   fill=BK)

    # ── Centre panel ─────────────────────────────────────────────
    orientation_diagrams(draw, accent, show_active_ls=True)

    # ── Right panel ──────────────────────────────────────────────
    # Stacked: STEP 1 (WiFi join QR, dynamic) above STEP 2 (URL QR, static
    # → http://192.168.4.1/, baked into bg). The URL QR is the trick that
    # forces iOS to open the captive portal: scanning a URL QR launches
    # Safari, Safari hits 192.168.4.1, iOS sees the request go to a
    # captive network and renders the portal instead of fighting whether
    # to auto-show CNA.
    cx = RIGHT_CX
    label_color = accent if accent != YL else BK

    # Step 1 — WiFi join (dynamic QR, overlaid by firmware)
    text_center(draw, cx, AP_QR_Y - 22, "STEP 1  —  JOIN WIFI", F_BIG, label_color)
    qx, qy, qp = AP_QR_X, AP_QR_Y, AP_QR_PX
    draw.rectangle([qx-6, qy-6, qx+qp+5, qy+qp+5], outline=accent, width=3)
    draw.rectangle([qx-3, qy-3, qx+qp+2, qy+qp+2], outline=BK, width=3)
    leave_qr_white(draw, qx, qy, qp)
    text_center(draw, cx, qy+qp+8, qr_label, F_FOOT, label_color)

    # Step 2 — URL QR (static, baked here so we don't need a second
    # firmware render path; URL is fixed at http://192.168.4.1/).
    url_qr = qrcode.QRCode(
        version=None,
        error_correction=qrcode.constants.ERROR_CORRECT_L,
        box_size=URL_QR_BOX,
        border=2,
    )
    url_qr.add_data("http://" + "192.168.4.1" + "/")
    url_qr.make(fit=True)
    url_img = url_qr.make_image(fill_color=BK, back_color=WH).convert("RGB")
    url_w, url_h = url_img.size
    url_x = cx - url_w // 2
    url_y = URL_QR_TARGET_Y

    text_center(draw, cx, url_y - 22, "STEP 2  —  OPEN PAGE", F_BIG, label_color)
    # Decorative border around the static URL QR — same accent treatment
    # as the WiFi QR for visual consistency.
    draw.rectangle([url_x-6, url_y-6, url_x+url_w+5, url_y+url_h+5], outline=accent, width=3)
    draw.rectangle([url_x-3, url_y-3, url_x+url_w+2, url_y+url_h+2], outline=BK, width=3)
    img.paste(url_img, (url_x, url_y))

    text_center(draw, cx, url_y + url_h + 8, "http://192.168.4.1/", F_FOOT, label_color)

    return img


def gen_ap_retry():
    """Step 1/2 with red accents + 'Connection Failed — try again' label,
    served after a failed WiFi connection attempt."""
    return gen_ap(
        accent=RD,
        header="CONNECTION FAILED  —  TRY AGAIN",
        qr_label="Connection Failed — try again",
    )


# ═══════════════════════════════════════════════════════════════════════════════
# SETUP SCREEN — green accent, account link
# ═══════════════════════════════════════════════════════════════════════════════
def gen_setup():
    img  = Image.new("RGB", (W, H), WH)
    draw = ImageDraw.Draw(img)

    # ── Status bar ────────────────────────────────────────────────
    draw.rectangle([0, 0, W-1, BAR_H-1], fill=GR)

    # WiFi bars icon
    bars = [(0, 8), (0, 13), (0, 18), (0, 22)]
    bx = 24
    for i, (_, bh) in enumerate(bars):
        draw.rectangle([bx + i*8, BAR_H//2 - bh//2, bx+i*8+5, BAR_H//2 + bh//2], fill=WH)
    draw.text((bx+38, 18), "WIFI CONNECTED  —  STEP 2 OF 2", font=F_BAR, fill=WH)

    # Right IP chip
    ip_text = "192.168.x.x"
    bb = draw.textbbox((0,0), ip_text, font=F_CHIP)
    chip_w = bb[2]-bb[0] + 22
    chip_x = W - chip_w - 20
    draw.rectangle([chip_x, 11, chip_x+chip_w, BAR_H-12], fill=WH)
    draw.text((chip_x+11, 18), ip_text, font=F_CHIP, fill=GR)

    # ── Panel dividers ────────────────────────────────────────────
    draw.rectangle([DIV1_X, BODY_Y, DIV1_X+1, H-1], fill=BK)
    draw.rectangle([DIV2_X, BODY_Y, DIV2_X+1, H-1], fill=BK)

    # ── Left panel ────────────────────────────────────────────────
    draw.text((28, BODY_Y+20), "Almost", font=F_HEAD, fill=BK)
    draw.text((28, BODY_Y+52), "ready.", font=F_HEAD, fill=BK)
    bb = draw.textbbox((0,0), "ready.", font=F_HEAD)
    draw.rectangle([28, BODY_Y+82, 28+bb[2]+2, BODY_Y+85], fill=GR)

    draw.text((28, BODY_Y+96), "Scan to name this frame and", font=F_STEP, fill=(80,80,75))
    draw.text((28, BODY_Y+110), "link it to your account.",    font=F_STEP, fill=(80,80,75))

    steps = [
        ("Scan the QR with your phone", "camera or QR app"),
        ("Sign in at pictureframe",     ".edholm.me"),
        ("Name the frame, choose",      "orientation — done."),
    ]
    sy = BODY_Y + 136
    for i, (l1, l2) in enumerate(steps):
        bx, by = 28, sy + i*46
        draw.rectangle([bx, by, bx+24, by+24], fill=GR)
        text_center(draw, bx+12, by+6, str(i+1), F_STEPN, WH)
        draw.text((62, by+3),  l1, font=F_STEP, fill=BK)
        draw.text((62, by+17), l2, font=F_STEP, fill=BK)

    # URL bar
    url_y = BODY_Y + 278
    draw.rectangle([28, url_y, 284, url_y+32], fill=BK)
    draw.text((38, url_y+4),  "URL", font=F_TINY, fill=GR)
    draw.text((38, url_y+16), "pictureframe.edholm.me/setup/...", font=ttf("DejaVuSans.ttf", 10), fill=WH)

    # Progress track
    prog_y = BODY_Y + 328
    draw.text((28, prog_y), "SETUP PROGRESS", font=F_TINY, fill=(140,140,135))
    seg_y = prog_y + 14
    segs = [("WiFi", GR), ("Account", BK), ("Frame ready", (200,200,195))]
    seg_w = (284 - 28 - 8) // 3  # ~82px each
    for i, (label, color) in enumerate(segs):
        sx = 28 + i*(seg_w+4)
        draw.rectangle([sx, seg_y, sx+seg_w, seg_y+6], fill=color)
        text_center(draw, sx+seg_w//2, seg_y+10, label, ttf("DejaVuSans.ttf", 9), BK)

    # ── Centre panel ─────────────────────────────────────────────
    orientation_diagrams(draw, GR, show_active_ls=True)

    # ── Right panel ──────────────────────────────────────────────
    cx = RIGHT_CX

    text_center(draw, cx, SETUP_QR_Y - 26, "SCAN TO FINISH", F_BIG, BK)

    qx, qy, qp = SETUP_QR_X, SETUP_QR_Y, SETUP_QR_PX
    draw.rectangle([qx-6, qy-6, qx+qp+5, qy+qp+5], outline=GR, width=3)
    draw.rectangle([qx-3, qy-3, qx+qp+2, qy+qp+2], outline=BK, width=3)

    leave_qr_white(draw, qx, qy, qp)

    # MAC chip below QR
    mac = "1C:C3:AB:D1:91:F8"
    bb = draw.textbbox((0,0), mac, font=F_CHIP)
    mw = bb[2]-bb[0]+20
    mx = cx - mw//2
    draw.rectangle([mx, qy+qp+8, mx+mw, qy+qp+26], fill=BK)
    draw.text((mx+10, qy+qp+11), mac, font=ttf("DejaVuSans-Bold.ttf", 10), fill=WH)

    return img


# ── Save ──────────────────────────────────────────────────────────────────────
def save_bin(img, path, preview_path):
    data = pack(img)
    with open(path, "wb") as f: f.write(data)
    print(f"Written {len(data):,} bytes → {os.path.abspath(path)}")

    # Reconstruct preview from packed data for verification
    prev = Image.new("RGB", (W, H))
    px = prev.load()
    for y in range(H):
        for x in range(0, W, 2):
            byte = data[y*(W//2) + x//2]
            px[x,   y] = PALETTE_RGB.get(byte >> 4,  (128,128,128))
            px[x+1, y] = PALETTE_RGB.get(byte & 0xF, (128,128,128))
    prev.save(preview_path)
    print(f"Preview → {os.path.abspath(preview_path)}")

if __name__ == "__main__":
    # Per-panel output directory: data/{vendor}-v{N}/. Defaults to the V1
    # 7.3" panel since that's the only one in production; pass --panel to
    # target a different one once new panels exist.
    panel = "waveshare73-v1"
    if "--panel" in sys.argv:
        panel = sys.argv[sys.argv.index("--panel") + 1]

    out_dir = os.path.join(os.path.dirname(__file__), f"../data/{panel}")
    os.makedirs(out_dir, exist_ok=True)

    print(f"Generating AP screen for {panel}…")
    save_bin(gen_ap(),       f"{out_dir}/ap_bg.bin",       f"{out_dir}/ap_bg_preview.png")
    print()
    print(f"Generating AP retry screen for {panel}…")
    save_bin(gen_ap_retry(), f"{out_dir}/ap_bg_retry.bin", f"{out_dir}/ap_bg_retry_preview.png")
    print()
    print(f"Generating setup screen for {panel}…")
    save_bin(gen_setup(),    f"{out_dir}/setup_bg.bin",    f"{out_dir}/setup_bg_preview.png")
    print()
    print("QR overlay constants for epd.cpp:")
    print(f"  AP_QR_X={AP_QR_X}, AP_QR_Y={AP_QR_Y}, AP_QR_CELL={AP_QR_CELL}, AP_QR_PX={AP_QR_PX}")
    print(f"  SETUP_QR_X={SETUP_QR_X}, SETUP_QR_Y={SETUP_QR_Y}, SETUP_QR_CELL={SETUP_QR_CELL}, SETUP_QR_PX={SETUP_QR_PX}")