#!/usr/bin/env python3
"""
Simulate the CodexBar ESP32 display and save a PNG for the README.
Run: python3 simulate_display.py
Out: docs/display_simulation.png
"""
import argparse
import math
from PIL import Image, ImageDraw, ImageFont
import os

from companion import pace as pace_lib

S = 3  # scale: device pixels → image pixels
DW, DH = 170, 320
W, H = DW * S, DH * S  # 510 × 960

# ── Colors — Ember palette from CodexBarRedesign/HANDOFF.md ───────────────
BG         = ( 12,  10,   8)  # #0c0a08
TRACK      = ( 51,  44,  35)  # #332c23
UNDER      = ( 84, 177, 115)  # #54b173
GOLD       = (214, 163,  65)  # #d6a341
ONPACE     = (220, 139,  52)  # #dc8b34
OVER1      = (200, 100,  58)  # #c8643a
OVER2      = (189,  67,  41)  # #bd4329
REDLINE    = (244, 238, 222)  # #f4eede
VALUE      = (243, 236, 220)  # #f3ecdc
VALUE_DIM  = (167, 160, 146)
GRAY       = (129, 119, 103)
DIV        = ( 34,  29,  23)
TANK_EMPTY = ( 23,  19,  15)
LIQUID_BLUE   = ( 79, 147, 201)  # #4f93c9 — actual fuel
LIQUID_GREEN  = ( 85, 181, 114)  # #55b572 — bonus above expected
LIQUID_DEFICIT = ( 82,  33,  25) # fDeficit blended over tank empty
ATTN_TEAL  = ( 94, 160, 147)
CLAUDE_C   = OVER1
CODEX_C    = (142, 154, 170)
WHITE      = VALUE

# ── Layout (image pixels = device pixels × S) ─────────────────────────────
HEADER_H = 2 * S    # 6px — minimal top edge
FOOTER_H = 2 * S    # 6px — minimal bottom edge
RO_LG, RI_LG = 65 * S, 46 * S   # large gauge (single provider)
RO_SM, RI_SM = 35 * S, 23 * S   # small gauge (two providers)

# ── Arc constants: clock degrees, 0°=top, increasing clockwise ─────────────
ARC_START_CLOCK = 213
ARC_SWEEP_CLOCK = 294
ARC_END_CLOCK   = ARC_START_CLOCK + ARC_SWEEP_CLOCK
REDLINE_CLOCK   = 360


def try_font(size):
    for p in [
        "/System/Library/Fonts/Helvetica.ttc",
        "/System/Library/Fonts/HelveticaNeue.ttc",
        "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
        "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
        "/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf",
    ]:
        if os.path.exists(p):
            try:
                return ImageFont.truetype(p, size)
            except Exception:
                pass
    return ImageFont.load_default()


_ICON_DIR = os.path.join(os.path.dirname(__file__), "assets", "icons")
_ICON_MAP  = {"Claude": "claude", "Codex": "openai", "Gemini": "gemini"}
_icon_cache: dict = {}

def load_icon(name: str, height_px: int):
    """Return an RGBA icon image scaled to height_px, or None if not found."""
    key = (name, height_px)
    if key not in _icon_cache:
        fname = _ICON_MAP.get(name)
        path  = os.path.join(_ICON_DIR, f"{fname}.png") if fname else None
        if path and os.path.exists(path):
            img = Image.open(path).convert("RGBA")
            w, h = img.size
            new_w = int(w * height_px / h)
            _icon_cache[key] = img.resize((new_w, height_px), Image.LANCZOS)
        else:
            _icon_cache[key] = None
    return _icon_cache[key]


def try_bold_font(size):
    # Try bold variants first, fall back to regular
    for p, idx in [
        ("/System/Library/Fonts/Helvetica.ttc", 1),
        ("/System/Library/Fonts/HelveticaNeue.ttc", 1),
        ("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 0),
        ("/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf", 0),
    ]:
        if os.path.exists(p):
            try:
                return ImageFont.truetype(p, size, index=idx)
            except Exception:
                pass
    return try_font(size)


F_PCT_LG      = try_bold_font(30)
F_PCT_SM      = try_bold_font(22)
F_PCT_100     = try_bold_font(26)
F_LABEL       = try_bold_font(22)
F_LABEL_OUTER = F_LABEL
F_NAME   = try_font(30)   # provider name
F_TIME   = try_font(28)   # time remaining
F_RESET  = try_font(18)   # compact reset countdown under a gauge
F_HDR    = try_font(20)   # header / footer
F_ATTN   = try_font(14)   # attention project label


def thick_arc(draw, cx, cy, ro, ri, a0, a1, color):
    """Draw a donut arc segment CW from a0 to a1 (PIL angles)."""
    if a1 < a0:
        a1 += 360
    width = max(1, int(ro - ri))
    mid = (ro + ri) / 2
    draw.arc([cx - mid, cy - mid, cx + mid, cy + mid],
             a0, a1, fill=color, width=width)


def clock_to_pil(clock):
    """Map handoff clock degrees (0=top, CW) to PIL arc degrees."""
    return (clock + 270) % 360


def thick_clock_arc(draw, cx, cy, ro, ri, start, end, color):
    thick_arc(draw, cx, cy, ro, ri, clock_to_pil(start), clock_to_pil(end), color)


def point_clock(cx, cy, r, clock):
    a = math.radians(clock)
    return (cx + r * math.sin(a), cy - r * math.cos(a))


def pace_to_angle(pace):
    pace = max(0.0, min(2.0, pace))
    if pace <= 1.0:
        return ARC_START_CLOCK + (REDLINE_CLOCK - ARC_START_CLOCK) * pace
    return REDLINE_CLOCK + (ARC_END_CLOCK - REDLINE_CLOCK) * (pace - 1.0)


def heat_color(pace):
    stops = [
        (0.00, UNDER),
        (0.55, UNDER),
        (0.90, GOLD),
        (1.00, ONPACE),
        (1.40, OVER1),
        (2.00, OVER2),
    ]
    pace = max(0.0, min(2.0, pace))
    for (x0, c0), (x1, c1) in zip(stops, stops[1:]):
        if x0 <= pace <= x1:
            return lerp_color(c0, c1, (pace - x0) / max(0.001, x1 - x0))
    return stops[-1][1]


def draw_liquid_fill(draw, cx, cy, r, fill_frac, color):
    """Fill a circle of radius r as a liquid-level indicator (0=empty, 1=full)."""
    fill_frac = max(0.0, min(1.0, fill_frac))
    if fill_frac < 0.01 or r <= 0:
        return
    if fill_frac >= 0.99:
        draw.ellipse([cx - r, cy - r, cx + r, cy + r], fill=color)
        return
    # Use polygon instead of PIL chord — chord has unreliable fill direction at
    # certain angles (fills the wrong segment or spills outside bounds).
    wy = cy + r * (1.0 - 2.0 * fill_frac)   # water line y (top=full, bottom=empty)
    hw = math.sqrt(max(0.0, r * r - (wy - cy) ** 2))
    # Chord endpoints: right and left on the water line
    a_right = math.degrees(math.atan2(wy - cy,  hw)) % 360  # right endpoint angle
    a_left  = math.degrees(math.atan2(wy - cy, -hw)) % 360  # left endpoint angle
    # Arc clockwise from right to left going through 90° (PIL bottom)
    if a_left < a_right:
        a_left += 360
    n = 36
    pts = []
    for i in range(n + 1):
        a = math.radians(a_right + (a_left - a_right) * i / n)
        pts.append((cx + r * math.cos(a), cy + r * math.sin(a)))
    draw.polygon(pts, fill=color)


def lerp_color(c0, c1, t):
    t = max(0.0, min(1.0, t))
    return tuple(int(c0[i] + t * (c1[i] - c0[i])) for i in range(3))


def draw_gradient_arc(draw, cx, cy, ro, ri, a0, a1, c0, c1):
    """Smooth gradient arc — one tiny segment per degree for a clean blend."""
    if a1 <= a0:
        return
    n = max(1, int(round(a1 - a0)))
    step = (a1 - a0) / n
    for i in range(n):
        t = (i + 0.5) / n
        thick_arc(draw, cx, cy, ro, ri, a0 + i * step, a0 + (i + 1) * step, lerp_color(c0, c1, t))


def draw_expected_line(draw, cx, cy, r, expected_frac, color):
    """Draw a horizontal line across the liquid circle at the expected burn-rate level."""
    expected_frac = max(0.0, min(1.0, expected_frac))
    ey = cy + r * (1 - 2 * expected_frac)
    if abs(ey - cy) > r:
        return
    hw = int(math.sqrt(max(0, r * r - (ey - cy) ** 2)))
    draw.line([(cx - hw, ey), (cx + hw, ey)], fill=color, width=S)


def draw_expected_dashed_line(draw, cx, cy, r, expected_frac, color, width):
    expected_frac = max(0.0, min(1.0, expected_frac))
    if expected_frac <= 0.01 or expected_frac >= 0.99:
        return
    ey = cy + r * (1 - 2 * expected_frac)
    if abs(ey - cy) > r:
        return
    hw = int(math.sqrt(max(0, r * r - (ey - cy) ** 2)))
    dash = max(3, int(r * 0.28))
    gap = max(2, int(r * 0.18))
    x = cx - hw
    while x < cx + hw:
        x2 = min(cx + hw, x + dash)
        draw.line([(x, ey), (x2, ey)], fill=color, width=width)
        x += dash + gap


def tc(draw, x, y, text, font, color):
    """Draw text centered at (x, y)."""
    bb = draw.textbbox((0, 0), text, font=font)
    draw.text((x - (bb[2] - bb[0]) // 2, y - (bb[3] - bb[1]) // 2),
              text, fill=color, font=font)


def tc_outlined(draw, x, y, text, font, color, shadow=(0, 0, 0), width=2):
    """Draw text centered at (x, y) with a black outline for contrast."""
    for dx in range(-width, width + 1):
        for dy in range(-width, width + 1):
            if dx != 0 or dy != 0:
                tc(draw, x + dx, y + dy, text, font, shadow)
    tc(draw, x, y, text, font, color)


def ellipsize(text, max_chars):
    return text if len(text) <= max_chars else text[: max(1, max_chars - 3)] + "..."


def fmt_time(sec):
    if sec <= 0:
        return "now"
    m = sec // 60; h = m // 60; d = h // 24
    if d >= 2: return f"{d}d {h % 24}h"
    if h >= 1: return f"{h}h {m % 60}m"
    return f"{m}m"


def draw_value(draw, cx, cy, remaining, size):
    num = str(int(round(remaining)))
    large = size >= 80 * S
    font = F_PCT_100 if remaining >= 100 else (F_PCT_LG if large else F_PCT_SM)
    pct_font = try_bold_font(max(10, int((font.size if hasattr(font, "size") else 20) * 0.52)))
    bb = draw.textbbox((0, 0), num, font=font)
    num_w = bb[2] - bb[0]
    tc_outlined(draw, cx, cy, num, font, VALUE, shadow=BG, width=2)
    draw.text((cx + num_w // 2 + 2 * S, cy - (font.size if hasattr(font, "size") else 20) // 2),
              "%", fill=VALUE_DIM, font=pct_font)


def draw_gauge(draw, cx, cy, size, pct, resets, dur):
    aw = size * 0.105
    r_arc = size / 2 - aw / 2 - size * 0.03
    ro = int(r_arc + aw / 2)
    ri = int(r_arc - aw / 2)
    r_fuel = int(r_arc - aw - size * 0.055)

    thick_clock_arc(draw, cx, cy, ro, ri, ARC_START_CLOCK, ARC_END_CLOCK, TRACK)

    remaining = 100 if pct == 0 else max(0, min(100, 100 - pct))
    elapsed_frac = max(0.0, min(1.0, pace_lib.elapsed_fraction(resets, dur))) if dur > 0 else 0.0
    exp_usage = elapsed_frac * 100.0
    if pct > 0 and exp_usage > 0.5:
        pace = max(0.0, min(2.0, pace_lib.pace_ratio(pct, elapsed_frac)))
    elif pct > 0:
        pace = 2.0
    else:
        pace = 0.0

    angle = pace_to_angle(pace)
    if angle > ARC_START_CLOCK + 0.5:
        thick_clock_arc(draw, cx, cy, ro, ri, ARC_START_CLOCK, angle, heat_color(pace))

    # Fuel tank.
    actual_frac = remaining / 100.0
    expected_frac = max(0.0, min(1.0, resets / dur)) if dur > 0 else actual_frac
    draw.ellipse([cx - r_fuel, cy - r_fuel, cx + r_fuel, cy + r_fuel], fill=TANK_EMPTY)
    if actual_frac >= expected_frac:
        draw_liquid_fill(draw, cx, cy, r_fuel, actual_frac, LIQUID_GREEN)
        draw_liquid_fill(draw, cx, cy, r_fuel, expected_frac, LIQUID_BLUE)
    else:
        draw_liquid_fill(draw, cx, cy, r_fuel, expected_frac, LIQUID_DEFICIT)
        draw_liquid_fill(draw, cx, cy, r_fuel, actual_frac, LIQUID_BLUE)
    draw_expected_dashed_line(draw, cx, cy, r_fuel, expected_frac, REDLINE, max(1, int(size * 0.022)))
    draw.ellipse([cx - r_fuel, cy - r_fuel, cx + r_fuel, cy + r_fuel], outline=(8, 6, 4), width=max(1, int(size * 0.01)))

    # Redline tick.
    mid = (ro + ri) / 2
    p1 = point_clock(cx, cy, mid - aw / 2 - size * 0.012, REDLINE_CLOCK)
    p2 = point_clock(cx, cy, mid + aw / 2 + size * 0.012, REDLINE_CLOCK)
    draw.line([p1, p2], fill=REDLINE, width=max(2, int(size * 0.022)))

    draw_value(draw, cx, cy, remaining, size)


def draw_section(img, draw, sy, sh, sw, name, wins, attention=None, pulse=True):
    brand = CLAUDE_C if name == "Claude" else CODEX_C
    sec_dw = sw / S
    sec_dh = sh / S
    portrait = sec_dh >= 120
    base_w = 170.0 if portrait else 320.0
    base_h = 158.0 if portrait else 170.0
    unit = min(sec_dw / base_w, sec_dh / base_h)
    x_pad = (sec_dw - base_w * unit) * 0.5
    y_pad = (sec_dh - base_h * unit) * 0.5

    def px(x): return int((x_pad + x * unit) * S)
    def py(y): return sy + int((y_pad + y * unit) * S)
    def ps(v): return int(v * unit * S)

    if portrait:
        big_size = ps(88)
        sm_size = ps(54)
        cx1, cy1 = px(61), py(67)
        cx2, cy2 = px(126), py(66)
        icon_cx, icon_cy = px(151), py(22)
        label1_y, label2_y = py(122), py(108)
        reset1_y, reset2_y = py(139), py(124)
        glyph_size = ps(22)
    else:
        big_size = ps(110)
        sm_size = ps(72)
        cx1, cy1 = px(112), py(76)
        cx2, cy2 = px(216), py(57)
        icon_cx, icon_cy = px(286), py(25)
        label1_y, label2_y = py(136), py(101)
        reset1_y, reset2_y = py(153), py(122)
        glyph_size = ps(34)

    icon_size = glyph_size
    icon = load_icon(name, icon_size)
    icon_h = icon.height if icon else 36
    icon_w = icon.width  if icon else 36
    icon_ix = icon_cx - icon_w // 2
    icon_iy = icon_cy - icon_h // 2

    # Provider icon (or fallback to text + brand dot)
    if icon is not None:
        img.paste(icon, (icon_ix, icon_iy), mask=icon)
    else:
        bb = draw.textbbox((0, 0), name, font=F_NAME)
        nw, nh = bb[2] - bb[0], bb[3] - bb[1]
        nx = icon_cx - nw // 2
        dot_x = nx - 8 * S
        dot_cy = icon_iy + nh // 2
        r = 4 * S
        draw.ellipse([dot_x - r, dot_cy - r, dot_x + r, dot_cy + r], fill=brand)
        draw.text((nx, icon_iy), name, fill=WHITE, font=F_NAME)

    attention = attention or {}
    state = attention.get("state", "idle")
    if state in ("needs_user", "running", "done"):
        accent = ATTN_TEAL if state in ("needs_user", "running") else GRAY
        rail_w = (2 * S if pulse else 1 * S) if state == "needs_user" else (1 * S if state == "running" else 0)
        if rail_w:
            draw.rectangle([0, sy + 4 * S, rail_w, sy + sh - 4 * S], fill=accent)
            # Soft low-contrast glow in the active section.
            draw.rectangle([rail_w, sy + 4 * S, min(sw, rail_w + 5 * S), sy + sh - 4 * S], fill=(10, 18, 16))
        if state == "needs_user" and attention.get("count", 0) > 1:
            label = f"{attention['count']} waiting"
        elif (state in ("needs_user", "running")) and attention.get("primary_project"):
            label = attention.get("primary_project")
        elif state == "running":
            label = "working"
        else:
            label = attention.get("primary_project") or ("waiting" if state == "needs_user" else "review")
        # Draw attention label at the top of this section, left-aligned after the rail.
        lbl_color = accent
        lbl_font = F_ATTN
        bb = draw.textbbox((0, 0), label, font=lbl_font)
        lbl_x = px(10)
        lbl_y = py(8)
        draw.rounded_rectangle(
            [lbl_x - 2 * S, lbl_y - 1 * S, lbl_x + (bb[2] - bb[0]) + 4 * S, lbl_y + (bb[3] - bb[1]) + 3 * S],
            radius=3 * S,
            fill=(17, 36, 32),
            outline=accent,
            width=1 * S,
        )
        draw.text((lbl_x, lbl_y), label, font=lbl_font, fill=lbl_color)

    if not wins:
        return

    # Gauge row.
    draw_gauge(draw, cx1, cy1, big_size, wins[0]["pct"], wins[0]["resets"], wins[0]["dur"])
    tc_outlined(draw, cx1, label1_y, wins[0]["lbl"], F_LABEL, VALUE, shadow=BG, width=1)
    if wins[0]["pct"] > 0:
        tc(draw, cx1, reset1_y, fmt_time(wins[0]["resets"]), F_RESET, GRAY)

    if len(wins) >= 2:
        draw_gauge(draw, cx2, cy2, sm_size, wins[1]["pct"], wins[1]["resets"], wins[1]["dur"])
        tc_outlined(draw, cx2, label2_y, wins[1]["lbl"], F_LABEL, VALUE, shadow=BG, width=1)
        if wins[1]["pct"] > 0:
            tc(draw, cx2, reset2_y, fmt_time(wins[1]["resets"]), F_RESET, GRAY)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--attention", action="store_true", help="render sample waiting-project state")
    parser.add_argument("--no-data", action="store_true", help="render the no-data state")
    parser.add_argument("--startup", action="store_true", help="render the startup grace state")
    parser.add_argument("--out", default="docs/display_simulation.png")
    args = parser.parse_args()

    # Sample data chosen to illustrate pace gauge concept:
    #   Claude 5h : 78% used, 2h left  → elapsed 3h/5h → expected 60% → pace 1.30 (RED, over pace)
    #   Claude 7d : 23% used, 5d left  → elapsed 2d/7d → expected 28.6% → pace 0.80 (GREEN)
    #   Codex  5h : 14% used, 1.5h left → elapsed 3.5h/5h → expected 70% → pace 0.20 (GREEN)
    #   Codex  7d : 90% used, 4d left  → elapsed 3d/7d → expected 42.9% → pace 2.0  (RED, maxed)
    providers = [
        {"name": "Claude", "wins": [
            {"lbl": "5h", "pct": 78, "resets": 7200,   "dur": 18000},
            {"lbl": "7d", "pct": 23, "resets": 432000,  "dur": 604800},
        ]},
        {"name": "Codex", "wins": [
            {"lbl": "5h", "pct":  0, "resets": 17820,  "dur": 18000},   # sub-1% → 0, empty
            {"lbl": "7d", "pct": 90, "resets": 345600,  "dur": 604800},
        ]},
    ]

    if args.attention:
        providers[1]["attention"] = {
            "state": "needs_user",
            "count": 1,
            "primary_project": "CodexBar-ESP32",
        }

    n = len(providers)
    # N=1: rotate to landscape (wider than tall) for the single-LLM layout
    if n == 1:
        img_w, img_h = DH * S, DW * S   # 960 × 510
    else:
        img_w, img_h = W, H              # 510 × 960 portrait

    img  = Image.new("RGB", (img_w, img_h), BG)
    draw = ImageDraw.Draw(img)

    ctr_x, ctr_y = img_w // 2, img_h // 2

    if args.startup:
        tc(draw, ctr_x, ctr_y - 36 * S, "Starting up", try_font(38), WHITE)
        tc(draw, ctr_x, ctr_y - 8 * S, "Waiting for first data", F_HDR, GRAY)
        tc(draw, ctr_x, ctr_y + 16 * S, "WiFi and server settling", F_TIME, GRAY)
        tc(draw, ctr_x, ctr_y + 40 * S, "Checking for 75s", F_TIME, GRAY)
        os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
        img.save(args.out)
        print(f"Saved {args.out}  ({img_w}×{img_h}px, {S}× device scale)")
        return

    if args.no_data:
        tc(draw, ctr_x, ctr_y - 36 * S, "No data", try_font(38), WHITE)
        tc(draw, ctr_x, ctr_y - 8 * S, "Server not found", F_HDR, GRAY)
        tc(draw, ctr_x, ctr_y + 16 * S, "Run server.py", F_TIME, GRAY)
        tc(draw, ctr_x, ctr_y + 36 * S, "Retrying every 30s", F_TIME, GRAY)
        tc(draw, ctr_x, ctr_y + 56 * S, "See companion README", F_LABEL, GRAY)
        os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
        img.save(args.out)
        print(f"Saved {args.out}  ({img_w}×{img_h}px, {S}× device scale)")
        return

    content_h = img_h - HEADER_H - FOOTER_H
    sec_h = content_h // n
    sw = img_w   # section width = full image width (sections stack vertically)

    for i, p in enumerate(providers):
        sy = HEADER_H + i * sec_h
        draw_section(img, draw, sy, sec_h, sw, p["name"], p["wins"],
                     attention=p.get("attention"))

    # Divider between providers
    for i in range(1, n):
        dy = HEADER_H + i * sec_h
        draw.line([(0, dy), (img_w, dy)], fill=DIV, width=1)

    os.makedirs("docs", exist_ok=True)
    out = args.out
    img.save(out)
    print(f"Saved {out}  ({img_w}×{img_h}px, {S}× device scale)")


if __name__ == "__main__":
    main()