"""Does a DESIGN.md make Claude Code use the brand's colors and type? Three arms per page:
- none: just the task;
- name: the task plus "in <Brand>'s visual style";
- file: the brand's DESIGN.md in the folder and "follow DESIGN.md"; the agent may only use the Read tool.
Each answer is one HTML file. We measure the share of its distinct hex colors that are in the brand's DESIGN.md
color tokens, whether the brand's primary color appears, and whether its font family is named.

  python3 design-md-test.py run <workdir> <folder of DESIGN.md files> && python3 design-md-test.py score <workdir> <same folder>
"""
import concurrent.futures as cf
import json
import re
import shutil
import statistics
import subprocess
import sys
import tempfile
from pathlib import Path

import yaml

REPS = 2
BRANDS = {"airbnb": "Airbnb", "stripe": "Stripe", "jobright": "Jobright"}
PAGES = {
    "pricing": "a pricing page with three plans, a monthly or yearly toggle and one highlighted plan",
    "signup": "a sign-up form with name, email, password, a terms checkbox and a submit button",
    "landing": "a landing page hero with a headline, a subhead, two buttons and a three-item feature grid",
}
ARMS = ["none", "name", "file"]
TAIL = " Reply with one complete HTML file with inline CSS in a single ```html block."


def prompt(arm, brand, page):
    base = f"Build {PAGES[page]}."
    if arm == "name":
        return f"Build {PAGES[page]} in {BRANDS[brand]}'s visual style." + TAIL
    if arm == "file":
        return f"Build {PAGES[page]}. Follow the design system in DESIGN.md in this folder." + TAIL
    return base + TAIL


def run_one(work, src, arm, brand, page, rep):
    out = work / "runs" / f"{arm}__{brand}__{page}__{rep}.json"
    if out.exists():
        return
    with tempfile.TemporaryDirectory() as tmp:
        if arm == "file":
            shutil.copy(Path(src) / f"{brand}.md", Path(tmp) / "DESIGN.md")
        r = subprocess.run(["claude", "-p", prompt(arm, brand, page), "--model", "haiku", "--setting-sources", "project",
                            "--output-format", "json", "--allowedTools", "Read",
                            "--disallowedTools", "Bash,Edit,Write,NotebookEdit,WebFetch,WebSearch,Glob,Grep"],
                           cwd=tmp, capture_output=True, text=True, timeout=600)
    out.write_text(r.stdout if r.stdout.strip().startswith("{") else json.dumps({"error": r.stderr[-500:]}))


def run(work, src):
    work = Path(work)
    (work / "runs").mkdir(parents=True, exist_ok=True)
    jobs = [(a, b, p, r) for r in range(REPS) for b in BRANDS for p in PAGES for a in ARMS]
    with cf.ThreadPoolExecutor(4) as pool:
        for f in cf.as_completed([pool.submit(run_one, work, src, *j) for j in jobs]):
            f.result()


def norm(h):
    h = h.lower()
    return "#" + "".join(c * 2 for c in h[1:]) if len(h) == 4 else h[:7]


def brand_tokens(src, brand):
    text = Path(src, f"{brand}.md").read_text()
    tree = {}
    m = re.match(r"---\n(.*?)\n---", text, re.S)
    for b in ([m.group(1)] if m else []) + re.findall(r"```ya?ml\n(.*?)```", text, re.S):
        try:
            d = yaml.safe_load(b)
            if isinstance(d, dict):
                tree.update(d)
        except yaml.YAMLError:
            pass
    colors = {norm(v) for v in (tree.get("colors") or {}).values() if isinstance(v, str) and re.fullmatch(r"#[0-9a-fA-F]{3,8}", v.strip())}
    fam = next((v.get("fontFamily") for v in (tree.get("typography") or {}).values() if isinstance(v, dict) and v.get("fontFamily")), "")
    first = re.sub(r"['\"]", "", fam.split(",")[0]).strip()
    return colors, norm(tree["colors"]["primary"]), first


def score(work, src):
    work = Path(work)
    rows = []
    for f in sorted((work / "runs").glob("*.json")):
        arm, brand, page, rep = f.stem.split("__")
        d = json.loads(f.read_text())
        if "result" not in d:
            continue
        html = "\n".join(re.findall(r"```html\n(.*?)```", d["result"], re.S)) or d["result"]
        used = {norm(h) for h in re.findall(r"#[0-9a-fA-F]{6}\b|#[0-9a-fA-F]{3}\b", html)}
        palette, primary, font = brand_tokens(src, brand)
        rows.append({"arm": arm, "brand": brand, "page": page, "rep": int(rep), "colorsUsed": len(used),
                     "paletteShare": round(len(used & palette) / len(used), 3) if used else 0,
                     "primaryUsed": primary in used, "fontNamed": font.lower() in html.lower(),
                     "readDesignMd": d.get("num_turns", 1) > 1, "outputTokens": (d.get("usage") or {}).get("output_tokens"),
                     "costUsd": d.get("total_cost_usd")})

    def summ(rs):
        if not rs:
            return {"n": 0}
        return {"n": len(rs), "medianPaletteShare": statistics.median(r["paletteShare"] for r in rs),
                "pctPrimaryUsed": round(100 * sum(r["primaryUsed"] for r in rs) / len(rs)),
                "pctFontNamed": round(100 * sum(r["fontNamed"] for r in rs) / len(rs)),
                "medianColorsUsed": statistics.median(r["colorsUsed"] for r in rs),
                "medianOutputTokens": statistics.median(r["outputTokens"] or 0 for r in rs),
                "totalCostUsd": round(sum(r["costUsd"] or 0 for r in rs), 4)}
    groups = {a: summ([r for r in rows if r["arm"] == a]) for a in ARMS}
    for b in BRANDS:
        for a in ARMS:
            groups[f"{b}, {a}"] = summ([r for r in rows if r["arm"] == a and r["brand"] == b])
    json.dump({"setup": {"claudeCode": "2.1.314", "model": "haiku", "reps": REPS, "brands": list(BRANDS), "pages": PAGES},
               "groups": groups, "rows": rows}, sys.stdout, indent=1)


if __name__ == "__main__":
    {"run": run, "score": score}[sys.argv[1]](*sys.argv[2:4])