#!/usr/bin/env python3
"""Generate the default pictureFrame registration screen.

Full-field baby blue background, "pictureFrame" title, centered QR code, and
registration instructions. Output is a 4bpp binary packed for the Waveshare
7.3" display (always 800×480 on the wire; portrait content is rotated 90° CW
into that format so the frame can be mounted either way).

Nibble map (verified on hardware):
  BLACK=0x0  WHITE=0x1  YELLOW=0x2  RED=0x3  BLUE=0x5  GREEN=0x6

Usage:
  python3 scripts/generate_default_image.py --orientation landscape
  python3 scripts/generate_default_image.py --orientation portrait
  python3 scripts/generate_default_image.py --mac AA:BB:CC:DD:EE:FF --orientation portrait
"""

import argparse
from pathlib import Path

import numpy as np
import qrcode
from PIL import Image, ImageDraw, ImageFont

EPD_W, EPD_H = 800, 480   # physical display dimensions, always
BABY_BLUE = (135, 206, 235)
BLACK     = (0, 0, 0)
BASE_URL  = "https://pictureframe.edholm.me"

FONT_BOLD  = "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf"
FONT_PLAIN = "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"

# Verified on Waveshare 7.3" hardware. Same map applies to 13.3" Spectra 6.
_PALETTE = [
    (0x0, (  0,   0,   0)),   # Black
    (0x1, (255, 255, 255)),   # White
    (0x2, (255, 220,   0)),   # Yellow
    (0x3, (220,  40,  40)),   # Red
    (0x5, ( 20,  80, 200)),   # Blue
    (0x6, ( 50, 160,  50)),   # Green
]
_CODES   = np.array([c   for c,   _ in _PALETTE], dtype=np.uint8)
_PAL_RGB = np.array([rgb for _, rgb in _PALETTE], dtype=np.float32)


def _floyd_steinberg(img_rgb: np.ndarray) -> np.ndarray:
    h, w = img_rgb.shape[:2]
    arr = img_rgb.astype(np.float32)
    out = np.zeros((h, w), dtype=np.uint8)
    for y in range(h):
        for x in range(w):
            px  = np.clip(arr[y, x], 0, 255)
            idx = int(np.argmin(((_PAL_RGB - px) ** 2).sum(axis=1)))
            out[y, x] = _CODES[idx]
            err = px - _PAL_RGB[idx]
            if x + 1 < w:
                arr[y, x + 1]     += err * (7 / 16)
            if y + 1 < h:
                if x > 0:
                    arr[y + 1, x - 1] += err * (3 / 16)
                arr[y + 1, x]         += err * (5 / 16)
                if x + 1 < w:
                    arr[y + 1, x + 1] += err * (1 / 16)
    return out


def _make_canvas_landscape(url: str) -> Image.Image:
    """800×480 — frame mounted horizontally."""
    W, H = 800, 480
    canvas = Image.new("RGB", (W, H), BABY_BLUE)
    draw   = ImageDraw.Draw(canvas)

    font_title = ImageFont.truetype(FONT_BOLD,  52)
    font_body  = ImageFont.truetype(FONT_PLAIN, 26)

    title = "pictureFrame"
    tw = draw.textlength(title, font=font_title)
    draw.text(((W - tw) / 2, 26), title, fill=BLACK, font=font_title)

    qr_size = 268
    qr_x, qr_y = (W - qr_size) // 2, 108
    _paste_qr(canvas, url, qr_x, qr_y, qr_size)

    for i, line in enumerate(["Scan to register your frame", url]):
        lw = draw.textlength(line, font=font_body)
        draw.text(((W - lw) / 2, qr_y + qr_size + 16 + i * 34), line, fill=BLACK, font=font_body)

    return canvas


def _make_canvas_portrait(url: str) -> Image.Image:
    """Design at 480×800 (portrait), rotate 90° CCW → 800×480 for the display.
    Hold the frame with the RIGHT edge up to read portrait content correctly."""
    W, H = 480, 800
    canvas = Image.new("RGB", (W, H), BABY_BLUE)
    draw   = ImageDraw.Draw(canvas)

    font_title = ImageFont.truetype(FONT_BOLD,  52)
    font_body  = ImageFont.truetype(FONT_PLAIN, 26)

    title = "pictureFrame"
    tw = draw.textlength(title, font=font_title)
    draw.text(((W - tw) / 2, 36), title, fill=BLACK, font=font_title)

    qr_size = 320
    qr_x, qr_y = (W - qr_size) // 2, 130
    _paste_qr(canvas, url, qr_x, qr_y, qr_size)

    for i, line in enumerate(["Scan to register your frame", url]):
        lw = draw.textlength(line, font=font_body)
        draw.text(((W - lw) / 2, qr_y + qr_size + 20 + i * 36), line, fill=BLACK, font=font_body)

    # ROTATE_90 (90° CCW): hold frame with RIGHT edge up to see content upright.
    # Combined framing guide uses the same direction — keep them in sync.
    return canvas.transpose(Image.ROTATE_90)


def _half_canvas(url: str, w: int, h: int, title_size: int, body_size: int, qr_size: int, qr_y: int) -> Image.Image:
    """Build a registration layout at arbitrary (w, h) — used for each half of the combined image."""
    canvas = Image.new("RGB", (w, h), BABY_BLUE)
    draw   = ImageDraw.Draw(canvas)
    font_title = ImageFont.truetype(FONT_BOLD,  title_size)
    font_body  = ImageFont.truetype(FONT_PLAIN, body_size)

    title = "pictureFrame"
    tw = draw.textlength(title, font=font_title)
    draw.text(((w - tw) / 2, 20), title, fill=BLACK, font=font_title)

    _paste_qr(canvas, url, (w - qr_size) // 2, qr_y, qr_size)

    for i, line in enumerate(["Scan to register your frame", url]):
        lw = draw.textlength(line, font=font_body)
        draw.text(((w - lw) / 2, qr_y + qr_size + 10 + i * (body_size + 6)), line, fill=BLACK, font=font_body)

    return canvas


def _make_canvas_combined(url: str) -> Image.Image:
    """800×480 — left half shows landscape content upright; right half shows portrait
    content physically rotated 90° CCW, so it appears sideways when the frame is flat.
    The user rotates the frame to see each half correctly — making orientation obvious."""
    DIVIDER = 2
    LEFT_W  = (EPD_W - DIVIDER) // 2   # 399
    RIGHT_W = EPD_W - LEFT_W - DIVIDER  # 399

    combined = Image.new("RGB", (EPD_W, EPD_H), BABY_BLUE)
    draw     = ImageDraw.Draw(combined)
    font_lbl = ImageFont.truetype(FONT_BOLD, 18)

    # ── Left half: upright landscape layout ──────────────────────────────────
    left = _half_canvas(url, w=LEFT_W, h=EPD_H, title_size=36, body_size=18, qr_size=210, qr_y=76)
    combined.paste(left, (0, 0))
    draw.text((6, EPD_H - 26), "LANDSCAPE", fill=BLACK, font=font_lbl)

    # ── Divider ───────────────────────────────────────────────────────────────
    draw.rectangle([LEFT_W, 0, LEFT_W + DIVIDER - 1, EPD_H - 1], fill=BLACK)

    # ── Right half: portrait content designed at 480×399, rotated 90° CCW ────
    # Designing at 480(h)×399(w) then rotating 90° CCW gives 399(w)×480(h) to
    # fill the right half.  The content appears sideways on the flat display —
    # rotate the frame 90° CCW (left edge up) to read this half correctly.
    portrait_src = _half_canvas(url, w=480, h=RIGHT_W,
                                 title_size=36, body_size=18, qr_size=210, qr_y=76)
    # ROTATE_90 = 90° CCW: 480×399 → 399×480
    portrait_half = portrait_src.transpose(Image.ROTATE_90)
    combined.paste(portrait_half, (LEFT_W + DIVIDER, 0))

    # Label drawn rotated — PIL doesn't rotate text natively, so draw on a
    # temp canvas then rotate and composite into bottom-right corner.
    lbl_text  = "PORTRAIT"
    lbl_font  = ImageFont.truetype(FONT_BOLD, 18)
    lbl_w     = int(draw.textlength(lbl_text, font=lbl_font))
    lbl_h     = 24
    lbl_img   = Image.new("RGB", (lbl_w + 4, lbl_h), BABY_BLUE)
    ImageDraw.Draw(lbl_img).text((2, 2), lbl_text, fill=BLACK, font=lbl_font)
    lbl_rot   = lbl_img.transpose(Image.ROTATE_90)   # rotate label to match content
    lx = LEFT_W + DIVIDER + 4
    ly = EPD_H - lbl_rot.height - 4
    combined.paste(lbl_rot, (lx, ly))

    return combined


def _paste_qr(canvas: Image.Image, url: str, x: int, y: int, size: int) -> None:
    qr = qrcode.QRCode(error_correction=qrcode.constants.ERROR_CORRECT_M, box_size=10, border=2)
    qr.add_data(url)
    qr.make(fit=True)
    qr_img = qr.make_image(fill_color=BLACK, back_color=BABY_BLUE).convert("RGB")
    canvas.paste(qr_img.resize((size, size), Image.NEAREST), (x, y))


def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--orientation", choices=["landscape", "portrait", "combined"], default="landscape")
    ap.add_argument("--mac", default=None, metavar="AA:BB:CC:DD:EE:FF")
    ap.add_argument("--out", default=None, help="Output .bin path (default: firmware/data/img_<orientation>.bin)")
    ap.add_argument("--preview", default=None, help="Preview PNG path (default: /tmp/pictureframe_<orientation>_preview.png)")
    args = ap.parse_args()

    url     = f"{BASE_URL}/setup/{args.mac}" if args.mac else BASE_URL
    out     = Path(args.out) if args.out else Path(f"firmware/data/img_{args.orientation}.bin")
    preview = Path(args.preview) if args.preview else Path(f"/tmp/pictureframe_{args.orientation}_preview.png")

    print(f"Orientation : {args.orientation}")
    print(f"QR URL      : {url}")

    makers = {"landscape": _make_canvas_landscape, "portrait": _make_canvas_portrait, "combined": _make_canvas_combined}
    canvas = makers[args.orientation](url)
    assert canvas.size == (EPD_W, EPD_H), f"canvas is {canvas.size}, expected ({EPD_W}, {EPD_H})"

    canvas.save(str(preview))
    print(f"Preview     : {preview}")

    print(f"Dithering {EPD_W}×{EPD_H} to 6-color palette…")
    codes  = _floyd_steinberg(np.array(canvas, dtype=np.float32))
    packed = ((codes[:, 0::2] << 4) | codes[:, 1::2]).astype(np.uint8).tobytes()
    assert len(packed) == EPD_H * (EPD_W // 2)

    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_bytes(packed)
    print(f"Written     : {len(packed):,} bytes → {out}")


if __name__ == "__main__":
    main()