"""Build the capability gallery: one screenshot per thing the harness can do.

Run it on a real desktop (or under Xvfb) and it writes ``gallery/<id>.png`` for every
capability registered in this file, plus ``gallery/index.json`` and a gallery page.

The screenshots are of the harness's *actual* behaviour, not mock-ups:

* **terminal shots** are captured stdout/stderr from a real command, rendered to an image
* **desktop shots** are real screen captures taken while the harness acts on a real
  X11/Windows desktop (element boxes, the resolved click point, and the app's own
  reaction are drawn or visible)

Nothing here is faked: if a command fails, its real failure is what gets pictured.

    python showcase.py                 # everything
    python showcase.py --only clarify  # one capability by id substring
    python showcase.py --list          # list ids without rendering
"""

from __future__ import annotations

import json
import os
import platform
import shutil
import subprocess
import sys
import textwrap
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Callable, Optional

ROOT = Path(__file__).resolve().parent
GALLERY = ROOT / "gallery"
PY = sys.executable
MONO_TTFS = (
    "/usr/share/fonts/truetype/dejavu/DejaVuSansMono.ttf",          # Debian/Ubuntu
    "C:/Windows/Fonts/consola.ttf",                                  # Windows (Consolas)
    "C:/Windows/Fonts/lucon.ttf",                                    # Windows (Lucida)
    "/System/Library/Fonts/Menlo.ttc",                               # macOS
    "/Library/Fonts/Menlo.ttc",
)
MONO_BOLD_TTFS = (
    "/usr/share/fonts/truetype/dejavu/DejaVuSansMono-Bold.ttf",
    "C:/Windows/Fonts/consolab.ttf",
    "/System/Library/Fonts/Menlo.ttc",
)
SCREEN_W, SCREEN_H = 1280, 800
TERM_COLS = 108
TERM_BG = (13, 17, 23)
TERM_FG = (201, 209, 217)
TERM_GREEN = (63, 185, 80)
TERM_RED = (248, 81, 73)
TERM_DIM = (139, 148, 158)

CATALOG: list["Shot"] = []


@dataclass
class Shot:
    """One capability, and the evidence that it works."""

    id: str
    category: str
    title: str
    shows: list[str]                 # the public callables this picture demonstrates
    note: str = ""                  # one line the card can print under the image
    kind: str = "terminal"          # "terminal" | "desktop"
    argv: list[str] | None = None   # for terminal shots: the command to run
    env: Optional[dict] = None      # extra environment for the command
    code: str | None = None         # for terminal shots: a python snippet to run
    scene: str | None = None        # for desktop shots: a named scene below
    make: Optional[Callable] = None  # for desktop shots: builds the picture
    prepare: Optional[Callable] = None  # for terminal shots: set the scene up first


def shot(**kwargs):
    def wrap(fn=None):
        s = Shot(**kwargs)
        CATALOG.append(s)
        return fn
    return wrap


# --------------------------------------------------------------------------- #
# rendering
# --------------------------------------------------------------------------- #

def _fonts(size=15, bold=False):
    from PIL import ImageFont

    for path in (MONO_BOLD_TTFS if bold else MONO_TTFS):
        try:
            return ImageFont.truetype(path, size)
        except Exception:
            continue
    return ImageFont.load_default()      # last resort: still produces a picture


def _wrap(text: str, cols: int) -> list[str]:
    out: list[str] = []
    for line in text.replace("\t", "    ").splitlines():
        if not line:
            out.append("")
            continue
        out.extend(textwrap.wrap(line, cols, drop_whitespace=False,
                                 replace_whitespace=False) or [""])
    return out


def terminal_image(title: str, output: str, subtitle: str = "") -> "Image.Image":
    """Render captured output as a terminal-window screenshot."""
    from PIL import Image, ImageDraw

    body_font, title_font = _fonts(15), _fonts(14, bold=True)
    lines = _wrap(output.rstrip("\n") or "(no output)", TERM_COLS)
    wrapped_title = _wrap(title, 78)
    sub = _wrap(subtitle, 96) if subtitle else []
    line_h, pad = 21, 18
    head_h = 46 + 20 * len(wrapped_title) + 19 * len(sub)
    height = head_h + line_h * len(lines) + 26

    image = Image.new("RGB", (1240, height), TERM_BG)
    draw = ImageDraw.Draw(image)
    draw.rectangle([0, 0, 1240, 34], fill=(22, 27, 34))
    for i, colour in enumerate([(255, 95, 86), (255, 189, 46), (39, 201, 63)]):
        draw.ellipse([18 + i * 22, 12, 30 + i * 22, 24], fill=colour)
    y = 46
    for line in wrapped_title:
        draw.text((24, y), line, font=title_font, fill=(230, 237, 243))
        y += 20
    for line in sub:
        draw.text((24, y), line, font=_fonts(13), fill=TERM_DIM)
        y += 19
    y = head_h
    for line in lines:
        colour = TERM_FG
        stripped = line.strip()
        if stripped.startswith("$ ") or stripped.startswith("python -m"):
            colour = TERM_GREEN
        elif any(word in stripped for word in ("result: failure", "ERROR", "FAILED",
                                               "refused", "Traceback", "Error:")):
            colour = TERM_RED
        elif stripped.startswith("result: success") or stripped.startswith("score: correct"):
            colour = TERM_GREEN
        draw.text((24, y), line, font=body_font, fill=colour)
        y += line_h
    return image


def run_argv(argv: list[str], *, extra_env: dict | None = None) -> tuple[int, str]:
    env = dict(os.environ, **(extra_env or {}))
    proc = subprocess.run(argv, capture_output=True, text=True, timeout=900, env=env)
    text = proc.stdout + (("\n" + proc.stderr) if proc.stderr.strip() else "")
    return proc.returncode, text


def _header(display: str) -> str:
    return (f"$ {display}" + chr(10))


# --------------------------------------------------------------------------- #
# the desktop: a real X server, real capture, real input
# --------------------------------------------------------------------------- #

DISPLAY_PROC = None


def _display_works() -> bool:
    return subprocess.run(["xdpyinfo"], capture_output=True).returncode == 0


def ensure_display() -> str:
    """Give the process a real display, starting Xvfb on Linux when there is none.

    Verified rather than assumed: a stale ``DISPLAY`` pointing at a dead server is
    worse than no display at all, because every scene then fails on ``tk.Tk()``.
    """
    global DISPLAY_PROC
    if platform.system() in ("Windows", "Darwin"):
        return os.environ.get("DISPLAY", "(native)")
    if os.environ.get("DISPLAY") and _display_works():
        return os.environ["DISPLAY"]
    if shutil.which("Xvfb") is None:
        return "(none)"
    os.environ["DISPLAY"] = ":99"
    Path.home().joinpath(".Xauthority").touch()
    os.environ["XAUTHORITY"] = str(Path.home() / ".Xauthority")
    DISPLAY_PROC = subprocess.Popen(
        ["Xvfb", ":99", "-screen", "0", f"{SCREEN_W}x{SCREEN_H}x24", "-ac", "-nolisten", "tcp"],
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, start_new_session=True,
    )
    for _ in range(80):
        if _display_works():
            break
        time.sleep(0.25)
    return os.environ["DISPLAY"]


SCENES: dict[str, str] = {}


def scene(name: str, source: str) -> str:
    """Register a Tk app used as the desktop a capability is demonstrated on."""
    SCENES[name] = source
    return name


def open_scene(name: str, tmp: Path) -> subprocess.Popen:
    """Launch a scene on the real display and wait until its window is mapped."""
    marker = tmp / f"{name}.marker"
    source = SCENES[name].replace("__MARKER__", str(marker))
    script = tmp / f"{name}.py"
    script.write_text(source)
    proc = subprocess.Popen([PY, str(script)], stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    ready = Path(str(marker) + ".ready")
    deadline = time.time() + 25
    while time.time() < deadline and not ready.exists():
        if proc.poll() is not None:
            raise RuntimeError(f"scene {name!r} exited early: {proc.stderr.read().decode()[:400]}")
        time.sleep(0.2)
    if not ready.exists():
        raise RuntimeError(f"scene {name!r} never mapped a window")
    time.sleep(0.6)                      # let the toolkit finish painting
    return proc


def close_scene(proc: subprocess.Popen) -> None:
    proc.terminate()
    try:
        proc.wait(timeout=10)
    except subprocess.TimeoutExpired:
        proc.kill()
        proc.wait(timeout=5)
    time.sleep(0.5)


def grab() -> "Image.Image":
    from tiny_computer_agent import GrabScreen
    return GrabScreen(monitor=1).capture().frame


def annotate(image, elements, *, click=None, title=""):
    """Draw what the harness *saw* (boxes + labels) and where it clicked."""
    from PIL import ImageDraw
    out = image.convert("RGB").copy()
    draw = ImageDraw.Draw(out)
    font = _fonts(13, bold=True)
    for e in elements:
        x, y, w, h = e.x, e.y, max(e.width, 1), max(e.height, 1)
        colour = (255, 138, 0) if e.kind != "grid" else (90, 120, 200)
        draw.rectangle([x, y, x + w, y + h], outline=colour, width=2)
        tag = f"{e.label}"[:34]
        tw = draw.textlength(tag, font=font)
        draw.rectangle([x, max(0, y - 17), x + tw + 8, y], fill=colour)
        draw.text((x + 4, max(0, y - 16)), tag, font=font, fill=(16, 16, 16))
    if click is not None:
        cx, cy = click
        draw.line([cx - 22, cy, cx + 22, cy], fill=(255, 0, 60), width=2)
        draw.line([cx, cy - 22, cx, cy + 22], fill=(255, 0, 60), width=2)
        draw.ellipse([cx - 9, cy - 9, cx + 9, cy + 9], outline=(255, 0, 60), width=3)
    if title:
        draw.rectangle([0, 0, out.width, 30], fill=(20, 22, 28))
        draw.text((12, 8), title, font=_fonts(15, bold=True), fill=(255, 255, 255))
    return out


def save(image, sid: str) -> Path:
    GALLERY.mkdir(exist_ok=True)
    target = GALLERY / f"{sid}.png"
    image.save(target, optimize=True)
    return target


# --------------------------------------------------------------------------- #
# rendering a shot
# --------------------------------------------------------------------------- #

def render_terminal(s: "Shot", tmp: Optional[Path] = None) -> "Image.Image":
    """Run the shot's command for real and picture what it printed."""
    cleanup = s.prepare(tmp) if s.prepare is not None else None
    try:
        return _terminal_body(s)
    finally:
        if cleanup is not None:
            cleanup()


def _terminal_body(s: "Shot") -> "Image.Image":
    if s.argv is not None:
        argv = [str(a).replace("{py}", PY) for a in s.argv]
        _, text = run_argv(argv, extra_env=s.env)
        shown = " ".join("python" if a == PY else str(a) for a in s.argv)
    else:
        argv = [PY, "-c", s.code]
        _, text = run_argv(argv, extra_env=s.env)
        shown = "python -c  # " + (s.note or s.title)
    return terminal_image(s.title, text, shown)


def render_desktop(s: "Shot", tmp: Path) -> "Image.Image":
    """Run the shot's desktop scene for real and picture the harness acting on it."""
    if s.make is None:
        raise RuntimeError(f"desktop shot {s.id!r} has no make() callable")
    if not os.environ.get("DISPLAY"):
        raise RuntimeError(f"desktop shot {s.id!r} needs a display")
    return s.make(tmp)


def overlay_lines(image, lines, *, title=""):
    """A dark strip of text under a picture - for the episode trace next to the pixels."""
    from PIL import ImageDraw
    out = image.convert("RGB").copy()
    draw = ImageDraw.Draw(out)
    font = _fonts(15)
    strip = 24 * len(lines) + 14
    draw.rectangle([0, out.height - strip, out.width, out.height], fill=(13, 17, 23))
    y = out.height - strip + 7
    for line in lines:
        draw.text((12, y), line, font=font, fill=TERM_FG)
        y += 24
    if title:
        draw.rectangle([0, 0, out.width, 30], fill=(20, 22, 28))
        draw.text((12, 8), title, font=_fonts(15, bold=True), fill=(255, 255, 255))
    return out


def side_by_side(left, right, *, title="", gap_label=("before", "after")):
    """Two frames of the same window, split down the middle."""
    from PIL import Image, ImageDraw
    width, height = left.width + right.width + 4, max(left.height, right.height) + 34
    out = Image.new("RGB", (width, height), (13, 17, 23))
    draw = ImageDraw.Draw(out)
    out.paste(left, (0, 34))
    out.paste(right, (left.width + 4, 34))
    draw.rectangle([left.width, 34, left.width + 4, height], fill=(255, 255, 255))
    draw.rectangle([0, 0, width, 30], fill=(20, 22, 28))
    draw.text((12, 8), title, font=_fonts(15, bold=True), fill=(255, 255, 255))
    draw.text((12, 16 + left.height), gap_label[0], font=_fonts(14, bold=True), fill=TERM_DIM)
    draw.text((left.width + 16, 16 + left.height), gap_label[1], font=_fonts(14, bold=True), fill=TERM_DIM)
    return out


def crop(image, box):
    return image.convert("RGB").crop(box)