import argparse
import sys
from pathlib import Path
from textwrap import wrap

try:
    import matplotlib.font_manager as font_manager
    import matplotlib.pyplot as plt
    from matplotlib.patches import FancyArrowPatch, FancyBboxPatch, Polygon, Rectangle
except ImportError as error:
    raise SystemExit(
        "matplotlib is required. Install it in your local Python environment, "
        "then rerun this script."
    ) from error


FONT_CANDIDATES = [
    "Microsoft YaHei",
    "SimHei",
    "Noto Sans CJK SC",
    "Source Han Sans SC",
    "PingFang SC",
    "Arial Unicode MS",
    "DejaVu Sans",
]


def configure_font():
    available = {font.name for font in font_manager.fontManager.ttflist}
    selected = next((font for font in FONT_CANDIDATES if font in available), "DejaVu Sans")
    plt.rcParams["font.sans-serif"] = [selected, "DejaVu Sans"]
    plt.rcParams["axes.unicode_minus"] = False


def wrap_mixed_text(text, width):
    lines = []
    for raw_line in text.splitlines():
        if not raw_line:
            lines.append("")
            continue
        if len(raw_line) <= width:
            lines.append(raw_line)
            continue
        if any(separator in raw_line for separator in ["_", "/", ".", ":"]) and " " not in raw_line:
            lines.append(raw_line)
            continue
        if " " in raw_line:
            lines.extend(wrap(raw_line, width=width, break_long_words=False, break_on_hyphens=False))
            continue
        lines.extend(raw_line[index : index + width] for index in range(0, len(raw_line), width))
    return "\n".join(lines)


def draw_lane(ax, y0, y1, label, color):
    ax.add_patch(
        Rectangle(
            (0.35, y0),
            17.3,
            y1 - y0,
            facecolor=color,
            edgecolor="none",
            alpha=0.34,
            zorder=0,
        )
    )
    ax.text(
        1.1,
        (y0 + y1) / 2,
        wrap_mixed_text(label, 8),
        fontsize=7.9,
        color="#475569",
        va="center",
        ha="center",
        weight="bold",
        linespacing=1.12,
    )


def add_box(ax, center, size, title, body, facecolor, edgecolor="#334155"):
    x, y = center
    width, height = size
    left = x - width / 2
    bottom = y - height / 2
    ax.add_patch(
        FancyBboxPatch(
            (left, bottom),
            width,
            height,
            boxstyle="round,pad=0.03,rounding_size=0.12",
            linewidth=1.3,
            facecolor=facecolor,
            edgecolor=edgecolor,
            zorder=2,
        )
    )
    ax.text(
        x,
        y + height * 0.22,
        wrap_mixed_text(title, 13),
        fontsize=9.7,
        weight="bold",
        color="#0f172a",
        ha="center",
        va="center",
        zorder=3,
    )
    ax.text(
        x,
        y - height * 0.17,
        wrap_mixed_text(body, 18),
        fontsize=8.0,
        color="#334155",
        ha="center",
        va="center",
        linespacing=1.12,
        zorder=3,
    )


def add_diamond(ax, center, size, title, body, facecolor, edgecolor="#334155"):
    x, y = center
    width, height = size
    points = [
        (x, y + height / 2),
        (x + width / 2, y),
        (x, y - height / 2),
        (x - width / 2, y),
    ]
    ax.add_patch(
        Polygon(points, closed=True, facecolor=facecolor, edgecolor=edgecolor, linewidth=1.3, zorder=2)
    )
    ax.text(
        x,
        y + 0.1,
        wrap_mixed_text(title, 12),
        fontsize=9.6,
        weight="bold",
        color="#0f172a",
        ha="center",
        va="center",
        zorder=3,
    )
    ax.text(
        x,
        y - 0.35,
        wrap_mixed_text(body, 14),
        fontsize=7.8,
        color="#334155",
        ha="center",
        va="center",
        linespacing=1.1,
        zorder=3,
    )


def add_arrow(ax, start, end, label=None, color="#475569", curve=0.0, dashed=False):
    arrow = FancyArrowPatch(
        start,
        end,
        arrowstyle="-|>",
        mutation_scale=15,
        linewidth=1.45,
        color=color,
        linestyle="--" if dashed else "-",
        connectionstyle=f"arc3,rad={curve}",
        zorder=1,
    )
    ax.add_patch(arrow)
    if label:
        mid_x = (start[0] + end[0]) / 2
        mid_y = (start[1] + end[1]) / 2
        ax.text(
            mid_x,
            mid_y + 0.22,
            label,
            fontsize=8,
            color=color,
            ha="center",
            va="center",
            bbox={"boxstyle": "round,pad=0.18", "facecolor": "#ffffff", "edgecolor": "none", "alpha": 0.9},
            zorder=4,
        )


def add_polyline_arrow(ax, points, label=None, color="#475569", dashed=False):
    if len(points) < 2:
        return
    for start, end in zip(points[:-2], points[1:-1]):
        ax.plot(
            [start[0], end[0]],
            [start[1], end[1]],
            color=color,
            linewidth=1.45,
            linestyle="--" if dashed else "-",
            zorder=1,
        )
    add_arrow(ax, points[-2], points[-1], color=color, dashed=dashed)
    if label:
        middle = points[len(points) // 2]
        ax.text(
            middle[0],
            middle[1] + 0.18,
            label,
            fontsize=8,
            color=color,
            ha="center",
            va="center",
            bbox={"boxstyle": "round,pad=0.18", "facecolor": "#ffffff", "edgecolor": "none", "alpha": 0.92},
            zorder=4,
        )


def draw_agent_loop(output_path, dpi):
    configure_font()
    fig, ax = plt.subplots(figsize=(18, 11), dpi=dpi)
    ax.set_xlim(0, 18)
    ax.set_ylim(0, 11)
    ax.axis("off")

    fig.patch.set_facecolor("#f8fafc")
    ax.set_facecolor("#f8fafc")

    draw_lane(ax, 8.65, 10.2, "用户聊天窗口\n只看聊天、上传、待确认命令", "#dbeafe")
    draw_lane(ax, 6.75, 8.25, "上下文与规格\nTauri / Rust 本地 Runtime", "#dcfce7")
    draw_lane(ax, 3.85, 6.35, "Agent Loop\n最多 3 轮，Evaluator 驱动返工", "#fef3c7")
    draw_lane(ax, 1.05, 3.45, "本地产物、预览\n和可审计证据", "#fce7f3")

    ax.text(
        8,
        10.72,
        "AI 游戏创作 App Agent Loop 结构",
        ha="center",
        va="center",
        fontsize=18,
        weight="bold",
        color="#0f172a",
    )
    ax.text(
        8,
        0.45,
        "事实源：memory/、assets/、game/、exports/、.agent/manifest.json、.agent/run.latest.json、.agent/passes/pass-N/",
        ha="center",
        va="center",
        fontsize=9,
        color="#64748b",
    )

    boxes = {
        "input": ((3.2, 9.42), (2.25, 0.92), "用户输入", "/project 授权\n普通需求触发生成"),
        "pending": ((5.9, 9.42), (2.25, 0.92), "Permission Gate", "pending 卡片\nconfirm / cancel 日志"),
        "tauri": ((8.6, 9.42), (2.35, 0.92), "Tauri Command", "generate_local_game_draft\n只操作授权路径"),
        "context": ((11.5, 9.42), (2.45, 0.92), "上下文装载", "LLM 配置\n记忆 + manifest\n资产摘要"),
        "planner": ((3.2, 7.5), (2.25, 0.92), "Planner", "整理规格\n写 .agent/spec.md"),
        "findings": ((5.9, 7.5), (2.25, 0.92), "Findings Seed", "初始化 Evaluator\n.agent/findings.md"),
        "orchestrator": ((8.6, 7.5), (2.35, 0.92), "Orchestrator", "agenda.md\ntask-graph.json"),
        "roles": ((11.5, 7.5), (2.45, 0.92), "6 组角色 Brief", "16 个组内任务\nactive / carry-over"),
        "generator": ((14.45, 7.5), (2.35, 0.92), "Generator", "结构化草案\nHTML + handoffs"),
        "repair": ((9.6, 5.0), (2.45, 1.0), "Repair Routes", "结构化问题路由\n扩展下游影响任务"),
        "failed": ((6.4, 5.0), (2.25, 1.0), "失败终止", "3 轮仍未通过\n只保留 trace 和 pass 快照"),
        "writer": ((12.7, 2.35), (2.45, 1.0), "Artifact Writer", "写 game/index.html\n设计、数值、资产\n发布草案"),
        "smoke": ((9.65, 2.35), (2.45, 1.0), "Playtest", "game.static_smoke\n拒绝占位\n拒绝危险 API"),
        "preview": ((6.6, 2.35), (2.4, 1.0), "本地预览", "127.0.0.1 HTTP\n外部浏览器打开"),
        "chat": ((3.55, 2.35), (2.25, 1.0), "聊天回显", "展示完成摘要\n/trace 读取 latest"),
        "trace": ((15.6, 2.35), (2.55, 1.18), "Trace / Evidence", "run.latest.json\nruns/ logs/ artifacts\nstep + toolCalls"),
    }

    colors = {
        "input": "#bfdbfe",
        "pending": "#c7d2fe",
        "tauri": "#bbf7d0",
        "context": "#bbf7d0",
        "planner": "#fde68a",
        "findings": "#fde68a",
        "orchestrator": "#fed7aa",
        "roles": "#fef08a",
        "generator": "#fde68a",
        "repair": "#fecaca",
        "failed": "#fee2e2",
        "writer": "#fbcfe8",
        "smoke": "#fbcfe8",
        "preview": "#fbcfe8",
        "chat": "#bfdbfe",
        "trace": "#e2e8f0",
    }

    for key, (center, size, title, body) in boxes.items():
        add_box(ax, center, size, title, body, colors[key])

    add_diamond(
        ax,
        (14.45, 5.0),
        (2.2, 1.48),
        "Evaluator",
        "质量门禁\n通过？",
        "#fed7aa",
    )

    add_arrow(ax, (4.33, 9.42), (4.77, 9.42))
    add_arrow(ax, (7.03, 9.42), (7.43, 9.42))
    add_arrow(ax, (9.78, 9.42), (10.28, 9.42))
    add_polyline_arrow(ax, [(11.5, 8.96), (11.5, 8.38), (3.2, 8.38), (3.2, 7.96)], label="开始 loop")
    add_arrow(ax, (4.33, 7.5), (4.77, 7.5))
    add_arrow(ax, (7.03, 7.5), (7.43, 7.5))
    add_arrow(ax, (9.78, 7.5), (10.27, 7.5))
    add_arrow(ax, (12.73, 7.5), (13.28, 7.5))
    add_polyline_arrow(ax, [(14.45, 7.04), (14.45, 6.03), (14.45, 5.74)])
    add_arrow(ax, (13.35, 5.0), (10.83, 5.0), label="不通过")
    add_polyline_arrow(ax, [(9.6, 5.5), (9.6, 6.15), (8.6, 6.15), (8.6, 7.04)], label="下一轮")
    add_arrow(ax, (8.37, 5.0), (7.57, 5.0), label="超过 3 轮")
    add_polyline_arrow(ax, [(14.45, 4.26), (14.45, 3.28), (12.7, 3.28), (12.7, 2.85)], label="通过")
    add_arrow(ax, (11.48, 2.35), (10.88, 2.35))
    add_arrow(ax, (8.43, 2.35), (7.8, 2.35))
    add_arrow(ax, (5.4, 2.35), (4.68, 2.35))
    add_arrow(ax, (13.93, 2.35), (14.33, 2.35), color="#64748b")
    add_polyline_arrow(
        ax,
        [(15.3, 7.08), (17.15, 7.08), (17.15, 2.95), (16.89, 2.7)],
        color="#64748b",
        dashed=True,
        label="持续写入 trace",
    )
    add_polyline_arrow(
        ax,
        [(15.6, 1.76), (15.6, 1.35), (3.55, 1.35), (3.55, 1.85)],
        color="#64748b",
        dashed=True,
    )

    output_path.parent.mkdir(parents=True, exist_ok=True)
    fig.savefig(output_path, bbox_inches="tight", facecolor=fig.get_facecolor())
    plt.close(fig)


def parse_args():
    parser = argparse.ArgumentParser(
        description="Draw the AI game creator Agent loop flowchart with matplotlib.",
    )
    parser.add_argument(
        "-o",
        "--output",
        default=".app/agent-loop-flow.png",
        help="Output image path. The extension controls the format, for example .png, .svg or .pdf.",
    )
    parser.add_argument("--dpi", type=int, default=180, help="Output DPI for raster formats.")
    return parser.parse_args()


def main():
    args = parse_args()
    output_path = Path(args.output)
    draw_agent_loop(output_path, args.dpi)
    print(f"agent loop flowchart written to {output_path.resolve()}")


if __name__ == "__main__":
    main()