#!/usr/bin/env python3
"""Computer-use benchmark: eight fixed tasks on a throwaway box, timed step by step.

Each task gets a fresh chat (the state file is reset from a template), runs through
`anybot --ask`, and is checked inside the box. Per task it records: success, steps, tool
calls by route (command / screen / browser / other), route switches, total seconds and
how they split between the model and the tools.

Never run it against the real box: it needs ANYBOT_BOX_NAME and ANYBOT_BOX_VOLUME of its
own (cloudbox.rs own_volume), and a data folder that is a copy, not the real one.

  python3 -I run.py --bin <anybot> --template <state.json> --data <dir> --docker-host <url> \
      --out results.json [--tasks T1,T4] [--runs 1]
"""
import argparse, json, os, re, shutil, subprocess, sys, time

BOX = os.environ.get("ANYBOT_BOX_NAME", "")
VOL = os.environ.get("ANYBOT_BOX_VOLUME", "")


def dx(cmd, docker_host, check=False, timeout=600, user="bot"):
    """Run a shell command in the bench box (as the bot user unless told otherwise)."""
    env = dict(os.environ, DOCKER_HOST=docker_host)
    r = subprocess.run(["docker", "exec", "-u", user, "-e", "DISPLAY=:1", "-e", "HOME=/home/bot", BOX, "bash", "-lc", cmd],
                       capture_output=True, text=True, env=env, timeout=timeout)
    if check and r.returncode != 0:
        raise RuntimeError(f"box command failed: {cmd}\n{r.stdout}\n{r.stderr}")
    return r.returncode, (r.stdout or "") + (r.stderr or "")


FORM = r"""<!doctype html><meta charset=utf-8><title>訂購表單</title>
<h1>訂購表單</h1>
<form id=f onsubmit="event.preventDefault();done()">
<label>姓名 <input name=name id=name></label><br>
<label>Email <input name=email id=email type=email></label><br>
<label>電話 <input name=phone id=phone></label><br>
<label>城市 <select name=city id=city><option value="">請選擇</option><option>台北</option><option>台中</option><option>高雄</option></select></label><br>
<label>數量 <input name=qty id=qty type=number min=1 max=9></label><br>
<label>備註 <textarea name=note id=note></textarea></label><br>
<label><input type=checkbox id=agree> 我同意條款</label><br>
<button type=submit>送出</button>
</form>
<p id=out></p>
<script>
function done(){
  const v=['name','email','phone','city','qty','note'].map(k=>document.getElementById(k).value.trim());
  if(!document.getElementById('agree').checked){document.getElementById('out').textContent='請先勾選同意';return;}
  let h=0; for(const c of v.join('|')){h=(h*31+c.codePointAt(0))>>>0;}
  document.getElementById('out').textContent='已送出，確認碼：'+(h%1000000).toString().padStart(6,'0');
}
</script>"""

FORM_VALUES = ["王小明", "ming@example.com", "0912345678", "台中", "3", "請在下午送達"]


def form_code():
    h = 0
    for ch in "|".join(FORM_VALUES):
        h = (h * 31 + ord(ch)) & 0xFFFFFFFF
    return f"{h % 1000000:06d}"


def digits(s):
    return re.sub(r"[,\s，]", "", s)


TASKS = [
    {"id": "T1", "route": "command", "name": "時區改成東京",
     # Earlier runs may have left a timezone in the bot's login files or the desktop's
     # environment (a workaround some runs used): clear every trace so `date` starts in UTC.
     "setup": "ln -sf /usr/share/zoneinfo/Etc/UTC /etc/localtime; rm -f /etc/timezone; "
              "for f in /home/bot/.profile /home/bot/.bashrc /home/bot/.bash_profile /home/bot/.pam_environment /etc/environment /etc/profile.d/*tz* /etc/profile.d/*time*; do "
              "[ -f \"$f\" ] && sed -i '/TZ=/d;/Asia\\/Tokyo/d' \"$f\"; done; "
              "grep -l 'TZ=' /proc/[0-9]*/environ 2>/dev/null | cut -d/ -f3 | xargs -r ps -o pid=,comm= -p | grep -E 'fluxbox|xterm' | awk '{print $1}' | xargs -r kill; sleep 1; "
              "su bot -c 'date +%Z'",
     "ask": "把你電腦的時區設成日本東京時間，設好後告訴我現在那台電腦顯示的時間。",
     "check": lambda dh, ans: "Tokyo" in dx("readlink -f /etc/localtime; cat /etc/timezone 2>/dev/null", dh)[1]},
    {"id": "T2", "route": "screen", "name": "計算機 App",
     "setup": "(command -v galculator >/dev/null || (sudo apt-get update -qq && sudo DEBIAN_FRONTEND=noninteractive apt-get install -y -qq galculator >/dev/null)); pkill galculator; true",
     "ask": "打開你電腦上的計算機 App（galculator），用它的按鈕算 1234 × 5678，告訴我計算機上顯示的答案。",
     "check": lambda dh, ans: "7006652" in digits(ans) and dx("pgrep -x galculator", dh)[0] == 0},
    {"id": "T3", "route": "screen-or-browser", "name": "Chromium 預設搜尋引擎",
     # Back to Chromium's own default (an earlier run may have set Bing): browser closed, the
     # saved choice removed from its preferences.
     "setup": "pkill -f '^/usr/lib/chromium/chromium' ; sleep 2; python3 - <<'PY'\n"
              "import json\nfor p in ['/home/bot/.config/anybot-profile/Default/Preferences', '/home/bot/.config/chromium/Default/Preferences']:\n"
              "    try:\n        d = json.load(open(p))\n    except Exception:\n        continue\n"
              "    d.pop('default_search_provider_data', None); d.pop('default_search_provider', None)\n    json.dump(d, open(p, 'w'))\n    print('reset', p)\nPY",
     "ask": "把你電腦上 Chromium 瀏覽器的預設搜尋引擎改成 Bing。",
     # Chromium writes its preferences up to ~10 s after a change: look for 15 s.
     "check": lambda dh, ans: "bing" in dx("for i in $(seq 15); do python3 - <<'PY' && break\nimport json, sys\np='/home/bot/.config/anybot-profile/Default/Preferences'\ntry:\n  t=json.dumps(json.load(open(p)).get('default_search_provider_data',{})).lower()\nexcept Exception as e: t=str(e)\nprint(t)\nsys.exit(0 if 'bing' in t else 1)\nPY\nsleep 1; done", dh)[1]},
    {"id": "T4", "route": "command", "name": "照片依月份分類",
     "setup": "rm -rf ~/bench/photos && mkdir -p ~/bench/photos && cd ~/bench/photos && "
              "for m in 2026-08 2026-09 2026-10; do for i in 1 2 3 4; do f=IMG_${m}_$i.jpg; echo x > $f; touch -d ${m}-1${i}T10:00:00 $f; done; done; "
              "for f in *; do mv \"$f\" \"photo_$RANDOM$RANDOM.jpg\"; done",
     "ask": "把 ~/bench/photos 裡的照片依照檔案的修改日期，分到用「年-月」命名的子資料夾（例如 2026-08）。",
     "check": lambda dh, ans: dx("cd ~/bench/photos && for d in 2026-08 2026-09 2026-10; do ls $d | wc -l; done | tr '\\n' ' '; ls -p | grep -v / | wc -l", dh)[1].split() == ["4", "4", "4", "0"]},
    {"id": "T5", "route": "command", "name": "建資料夾再改名",
     "setup": "mkdir -p ~/bench && rm -rf ~/bench/報告 ~/bench/'2026 報告'",
     "ask": "在 ~/bench 建一個叫「報告」的資料夾，再把它改名成「2026 報告」。",
     "check": lambda dh, ans: dx("test -d ~/bench/'2026 報告' && ! test -e ~/bench/報告 && echo OK", dh)[1].strip() == "OK"},
    {"id": "T6", "route": "command", "name": "裝 tree 查版本",
     "setup": "sudo apt-get remove -y -qq tree >/dev/null 2>&1; true",
     "ask": "在你的電腦上安裝 tree 這個指令，告訴我它的版本號。",
     "check": lambda dh, ans: dx("tree --version", dh)[0] == 0},
    {"id": "T7", "route": "browser", "name": "維基百科查高度",
     "setup": "true",
     "ask": "用你電腦上的瀏覽器到中文維基百科搜尋「台北101」，告訴我它的高度是多少公尺。",
     "check": lambda dh, ans: "508" in ans},
    {"id": "T8", "route": "browser", "name": "填表送出",
     "setup": "mkdir -p ~/bench/form && cat > ~/bench/form/index.html <<'EOF'\n" + FORM + "\nEOF\n"
              # (Anchored at the program: this very shell's command line names the server too.)
              "pkill -f '^python3 -m http.server 8770'; sleep 0.5; cd /home/bot/bench/form && (su bot -c 'setsid python3 -m http.server 8770 --bind 127.0.0.1 >/dev/null 2>&1 &'); "
              "for i in 1 2 3 4 5 6 7 8 9 10; do curl -s -o /dev/null http://127.0.0.1:8770/ && break; sleep 0.5; done; curl -s http://127.0.0.1:8770/ | grep -c 訂購表單",
     "ask": "到 http://127.0.0.1:8770 填表：姓名 王小明、Email ming@example.com、電話 0912345678、城市 台中、數量 3、備註 請在下午送達，勾選同意後送出，告訴我頁面上的確認碼。",
     "check": lambda dh, ans: form_code() in ans},
]

ROUTE = [("command", ("shell", "read_file", "write_file", "edit_file", "list_dir", "search", "read_terminal", "write_terminal")),
         ("screen", ("computer_use_",)),
         ("browser", ("browser_", "web_fetch"))]


def route_of(name):
    for r, prefixes in ROUTE:
        if any(name == p or name.startswith(p) for p in prefixes):
            return r
    return "other"


def analyze(msgs):
    """Steps, calls by route, switches, and time split from the turn's messages."""
    if not msgs:
        return {}
    t0 = msgs[0]["at"]
    prev = t0
    model_ms = tool_ms = 0
    calls, seq = [], []
    steps = 0
    for m in msgs[1:]:
        gap = max(0, m.get("at", prev) - prev)
        prev = m.get("at", prev)
        if m.get("role") == "Assistant":
            if m.get("summary"):
                continue
            steps += 1
            model_ms += gap
            for c in m.get("tool_calls") or []:
                calls.append(c["name"])
                seq.append(route_of(c["name"]))
        elif m.get("role") == "Tool":
            tool_ms += gap
    switches = sum(1 for a, b in zip(seq, seq[1:]) if a != b and "other" not in (a, b))
    by_route = {}
    for r in seq:
        by_route[r] = by_route.get(r, 0) + 1
    final = next((m.get("text", "") for m in reversed(msgs) if m.get("role") == "Assistant" and m.get("text")), "")
    # Every step, short, for reading why a run went the way it did.
    transcript = []
    for m in msgs:
        at = round((m.get("at", t0) - t0) / 1000, 1)
        if m.get("role") == "Assistant":
            transcript.append({"t": at, "say": (m.get("text") or "")[:300], "think": len(m.get("thinking") or ""),
                               "calls": [{"name": c["name"], "args": (c.get("args") or "")[:400]} for c in m.get("tool_calls") or []]})
        elif m.get("role") == "Tool":
            transcript.append({"t": at, "tool": m.get("tool_name"), "out": (m.get("text") or "")[:400]})
    return {"steps": steps, "calls": len(calls), "by_route": by_route, "switches": switches,
            "total_s": round((prev - t0) / 1000, 1), "model_s": round(model_ms / 1000, 1),
            "tool_s": round(tool_ms / 1000, 1), "tools": calls, "answer": final[:600], "transcript": transcript}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--bin", required=True)
    ap.add_argument("--template", required=True)
    ap.add_argument("--data", required=True)
    ap.add_argument("--docker-host", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--tasks", default="")
    ap.add_argument("--runs", type=int, default=1)
    ap.add_argument("--secs", type=int, default=600)
    ap.add_argument("--keep-skills", action="store_true", help="carry skills learned in one run into the next (to measure replay)")
    a = ap.parse_args()
    if not BOX or not VOL or BOX == "anybot-box" or VOL == "anybot-box-home":
        sys.exit("set ANYBOT_BOX_NAME and ANYBOT_BOX_VOLUME to a throwaway box of its own")
    real = os.path.realpath(os.path.expanduser("~/Library/Application Support/AnyBot"))
    if os.path.realpath(a.data).startswith(real):
        sys.exit("--data must be a scratch copy, not the real data folder")
    want = {t.strip() for t in a.tasks.split(",") if t.strip()}
    tasks = [t for t in TASKS if not want or t["id"] in want]
    results = json.load(open(a.out)) if os.path.exists(a.out) else []
    env = dict(os.environ, ANYBOT_DATA_DIR=a.data, ANYBOT_ASK_SECS=str(a.secs))
    kept_skills = []
    for run in range(a.runs):
        for t in tasks:
            print(f"== {t['id']} {t['name']} (run {run + 1})", flush=True)
            # Setup runs as the box's root (sudo can't raise the bot user under gVisor); what it
            # makes in the bot's home is handed back to the bot.
            # Every task starts on a clear desktop: apps a previous task left open are closed.
            dx("pkill -x galculator; pkill -x xterm; pkill -f '^/usr/lib/chromium/chromium'; sleep 1; true", a.docker_host, user="root")
            rc, out = dx(t["setup"].replace("sudo ", "") + "; chown -R bot:bot /home/bot/bench 2>/dev/null; true", a.docker_host, check=False, timeout=900, user="root")
            print("   setup:", out.strip().replace("\n", " ")[-160:], flush=True)
            shutil.copyfile(a.template, os.path.join(a.data, "state.json"))
            os.chmod(os.path.join(a.data, "state.json"), 0o600)
            if a.keep_skills and kept_skills:
                st0 = json.load(open(os.path.join(a.data, "state.json")))
                st0["skills"] = kept_skills
                json.dump(st0, open(os.path.join(a.data, "state.json"), "w"), ensure_ascii=False)
            started = time.time()
            try:
                p = subprocess.run([a.bin, "--ask", t["ask"]], capture_output=True, text=True, env=env, timeout=a.secs + 400)
                log = p.stdout + p.stderr
            except subprocess.TimeoutExpired as e:
                log = f"TIMEOUT {e}"
            wall = round(time.time() - started, 1)
            st = json.load(open(os.path.join(a.data, "state.json")))
            if a.keep_skills:
                kept_skills = st.get("skills", [])
            msgs = st["bots"][0]["messages"]
            r = analyze(msgs)
            try:
                ok = bool(t["check"](a.docker_host, r.get("answer", "")))
            except Exception as e:  # a checker that fails counts as a miss
                ok = False
                r["check_error"] = str(e)
            r["skills"] = [{"name": k.get("name"), "steps": len(k.get("steps", []))} for k in st.get("skills", [])]
            r["replayed"] = "replay_skill" in r.get("tools", [])
            r.update({"id": t["id"], "name": t["name"], "expected_route": t["route"], "ok": ok, "run": run + 1,
                      "process_s": wall, "at": time.strftime("%Y-%m-%d %H:%M"), "log_tail": log[-1500:]})
            results.append(r)
            json.dump(results, open(a.out, "w"), ensure_ascii=False, indent=1)
            print(f"   ok={ok} steps={r.get('steps')} total={r.get('total_s')}s model={r.get('model_s')}s tool={r.get('tool_s')}s "
                  f"routes={r.get('by_route')} switches={r.get('switches')}", flush=True)


if __name__ == "__main__":
    main()