"""Conceptual/architecture figures for the paper (no model server needed).

Produces two vector figures under paper/figures/:
  * architecture.pdf/png  -- the three memory arms feeding a draft-then-verify
                             speculative decoder (no-memory vs frozen global
                             datastore vs our persistent per-user evicting store).
  * safety_gate.pdf/png   -- the speculative-execution decision flow contrasting
                             naive-exec, confidence-gate, and idempotency-gate.
"""
from __future__ import annotations

from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.patches import FancyBboxPatch, FancyArrowPatch

ROOT = Path(__file__).resolve().parent.parent
FIGS = ROOT / "figures"

BLUE = "#1b6ca8"
RED = "#d1495b"
GREY = "#888888"
GREEN = "#2e8b57"
DARK = "#2b2b2b"


def _box(ax, x, y, w, h, text, fc="white", ec=DARK, fs=8.5, lw=1.3, tc=DARK,
         style="round,pad=0.02,rounding_size=0.06", bold=False):
    p = FancyBboxPatch((x, y), w, h, boxstyle=style, fc=fc, ec=ec, lw=lw,
                       mutation_aspect=1.0, zorder=2)
    ax.add_patch(p)
    ax.text(x + w / 2, y + h / 2, text, ha="center", va="center", fontsize=fs,
            color=tc, zorder=3, fontweight="bold" if bold else "normal",
            wrap=True)
    return p


def _arrow(ax, xy1, xy2, color=DARK, lw=1.4, style="-|>", ls="-", rad=0.0):
    a = FancyArrowPatch(xy1, xy2, arrowstyle=style, mutation_scale=12,
                        color=color, lw=lw, linestyle=ls, zorder=1,
                        connectionstyle=f"arc3,rad={rad}")
    ax.add_patch(a)


def architecture():
    fig, ax = plt.subplots(figsize=(10.5, 4.6))
    ax.set_xlim(0, 10.5)
    ax.set_ylim(0, 4.6)
    ax.axis("off")

    # user query stream (sessions)
    _box(ax, 0.15, 2.0, 1.5, 0.8,
         "User $u$ query\n(session $s$)", fc="#eef3f7", fs=8.5)
    ax.text(0.9, 1.75, "recurring, evolving\nper-user tool use",
            ha="center", va="top", fontsize=6.8, color=GREY, style="italic")

    # three arms as stacked draft sources
    arm_x = 2.35
    arms = [
        (3.55, GREY, "No memory",
         "schema-only draft\n(zero history)"),
        (2.30, RED, "Static datastore (ToolSpec-style)",
         "one global store,\nbuilt once then FROZEN"),
        (1.05, BLUE, "Ours: persistent per-user memory",
         "grows across sessions,\nLRU/LFU eviction,\npersonalized retrieval"),
    ]
    for y, color, title, sub in arms:
        _box(ax, arm_x, y, 3.0, 1.05, "", ec=color, lw=1.6,
             fc=color + "14" if False else "white")
        ax.text(arm_x + 0.12, y + 0.86, title, ha="left", va="center",
                fontsize=8.2, color=color, fontweight="bold")
        ax.text(arm_x + 0.12, y + 0.36, sub, ha="left", va="center",
                fontsize=7.2, color=DARK)
        _arrow(ax, (1.65, 2.4), (arm_x, y + 0.55), color=color, lw=1.2,
               rad=0.05)

    # draft -> speculative decoder
    dec_x = 6.15
    _box(ax, dec_x, 1.75, 1.95, 1.3,
         "Draft-then-verify\nspeculative\ndecoder", fc="#fff7e6",
         ec="#c8860a", fs=8.5, bold=True)
    for y, color, *_ in arms:
        _arrow(ax, (arm_x + 3.0, y + 0.55), (dec_x, 2.4), color=color, lw=1.1,
               ls="--", rad=-0.05)
    ax.text(dec_x + 0.98, 3.12, "draft tool call", ha="center", va="bottom",
            fontsize=7, color="#c8860a")

    # served target model
    _box(ax, dec_x - 0.1, 0.35, 2.15, 0.85,
         "Served target model\n(gpt-oss-120b, greedy)", fc="#f2f2f2", fs=7.6)
    _arrow(ax, (dec_x + 0.95, 1.75), (dec_x + 0.95, 1.2), color=DARK, lw=1.2)
    _arrow(ax, (dec_x + 1.15, 1.2), (dec_x + 1.15, 1.75), color=DARK, lw=1.2)
    ax.text(dec_x + 1.35, 1.47, "verify\n(token LCP)", ha="left", va="center",
            fontsize=6.6, color=DARK)

    # accepted output + metric
    _box(ax, 8.55, 1.75, 1.8, 1.3,
         "Accepted tokens\n= mean accepted\ntokens (MAT)", fc="#eaf5ee",
         ec=GREEN, fs=8.2, tc=DARK)
    _arrow(ax, (dec_x + 1.95, 2.4), (8.55, 2.4), color=GREEN, lw=1.6)

    # observe/write-back loop (target grows the personal store)
    _arrow(ax, (dec_x + 0.05, 1.9), (arm_x + 3.0, 1.5), color=BLUE, lw=1.1,
           ls=":", rad=0.28)
    ax.text(5.1, 1.05, "observe target $\\rightarrow$ write back\n"
            "(personal store only; static store frozen after warmup)",
            ha="center", va="center", fontsize=6.8, color=BLUE, style="italic")

    ax.set_title("Persistent, personalized tool-call memory feeding a "
                 "speculative decoder", fontsize=10.5, fontweight="bold")
    fig.tight_layout()
    FIGS.mkdir(parents=True, exist_ok=True)
    fig.savefig(FIGS / "architecture.pdf")
    fig.savefig(FIGS / "architecture.png", dpi=150)
    plt.close(fig)
    print("wrote architecture.{pdf,png}")


def safety_gate():
    fig, ax = plt.subplots(figsize=(9.5, 4.2))
    ax.set_xlim(0, 9.5)
    ax.set_ylim(0, 4.2)
    ax.axis("off")

    _box(ax, 0.2, 1.75, 1.75, 0.9,
         "Memory proposes\na drafted\ntool call", fc="#eef3f7", fs=8.2)

    # decision diamond-ish: idempotent?
    _box(ax, 2.35, 1.75, 1.6, 0.9, "Tool\nidempotent?", fc="#fff7e6",
         ec="#c8860a", fs=8.4, bold=True)
    _arrow(ax, (1.95, 2.2), (2.35, 2.2))

    # idempotent -> always speculatively execute (safe)
    _box(ax, 4.35, 3.0, 2.4, 0.85,
         "Speculatively EXECUTE\n(wrong $\\Rightarrow$ just re-run)",
         fc="#eaf5ee", ec=GREEN, fs=8.0)
    _arrow(ax, (3.95, 2.5), (4.35, 3.35), color=GREEN, rad=0.15)
    ax.text(4.15, 3.05, "yes", fontsize=7.5, color=GREEN, ha="left")

    # non-idempotent -> policy fork
    ax.text(4.15, 1.55, "no", fontsize=7.5, color=RED, ha="left")
    _box(ax, 4.35, 0.35, 2.4, 2.2, "", ec=RED, lw=1.4, fc="#fdeef1")
    ax.text(5.55, 2.35, "Non-idempotent tool\n(policy decides)",
            ha="center", va="center", fontsize=8.0, color=RED,
            fontweight="bold")
    _arrow(ax, (3.95, 1.95), (4.35, 1.6), color=RED, rad=-0.12)

    ax.text(4.5, 1.75, "naive:", fontsize=7.6, color=DARK, ha="left",
            fontweight="bold")
    ax.text(4.95, 1.75, "execute anyway", fontsize=7.6, color=RED, ha="left")
    ax.text(4.5, 1.32, "conf-gate:", fontsize=7.6, color=DARK, ha="left",
            fontweight="bold")
    ax.text(5.15, 1.32, "execute if conf$\\geq\\tau$", fontsize=7.6,
            color="#c8860a", ha="left")
    ax.text(4.5, 0.89, "ours:", fontsize=7.6, color=DARK, ha="left",
            fontweight="bold")
    ax.text(4.9, 0.89, "NEVER execute;", fontsize=7.6, color=GREEN, ha="left")
    ax.text(4.9, 0.58, "draft only, wait for verify", fontsize=7.0,
            color=GREEN, ha="left")

    # outcomes
    _box(ax, 7.15, 2.55, 2.15, 0.95,
         "Irreversible harm\nif draft wrong\n(severity-weighted cost)",
         fc="#fdeef1", ec=RED, fs=7.6, tc=RED)
    _box(ax, 7.15, 0.65, 2.15, 0.95,
         "Zero cost, keeps\nlatency win on\nsafe tools", fc="#eaf5ee",
         ec=GREEN, fs=7.8, tc=DARK)
    _arrow(ax, (6.75, 1.65), (7.15, 2.9), color=RED, rad=0.1, lw=1.2)
    _arrow(ax, (6.75, 0.95), (7.15, 1.1), color=GREEN, rad=-0.05, lw=1.4)

    ax.set_title("Idempotency-gated speculative execution: naive and "
                 "confidence gates still leak irreversible cost",
                 fontsize=10.0, fontweight="bold")
    fig.tight_layout()
    FIGS.mkdir(parents=True, exist_ok=True)
    fig.savefig(FIGS / "safety_gate.pdf")
    fig.savefig(FIGS / "safety_gate.png", dpi=150)
    plt.close(fig)
    print("wrote safety_gate.{pdf,png}")


if __name__ == "__main__":
    architecture()
    safety_gate()