"""CUA Tutor Demo — PC agent. Shows the consent modal; executes ONLY
manifest-allowed actions; Ctrl+Alt+F12 kills the session; pyautogui
failsafe (mouse to top-left corner) also aborts instantly.
Deps: pip install requests pyautogui pillow pynput
"""
import argparse, base64, io, sys, threading, time
import requests

p = argparse.ArgumentParser()
p.add_argument("--url", required=True)
p.add_argument("--sid", required=True)
p.add_argument("--key", required=True)
A = p.parse_args()
BASE = A.url.rstrip("/")
K = {"key": A.key}
URL = f"{BASE}/api/agent/{A.sid}"

STATE = {"kill": False}


def die(why):
    print(f"[killed] {why}", flush=True)
    STATE["kill"] = True
    try:
        requests.post(f"{BASE}/api/session/{A.sid}/kill", params=K, timeout=5)
    except Exception:
        pass


try:
    import pyautogui
    pyautogui.FAILSAFE = True
    pyautogui.PAUSE = 0.15
except Exception as e:
    print(f"pyautogui unavailable: {e}\nInstall: pip install pyautogui pillow requests pynput", flush=True)
    sys.exit(1)

# ---- kill hotkey ----
try:
    from pynput import keyboard
    from pynput.keyboard import Key as Kk
    _L = set()

    def on_press(k):
        _L.add(k)
        if Kk.f12 in _L and (Kk.ctrl_l in _L or Kk.ctrl_r in _L) \
                and (Kk.alt_l in _L or Kk.alt_r in _L):
            die("hotkey Ctrl+Alt+F12")
        return True

    def on_release(k):
        _L.discard(k)

    keyboard.Listener(on_press=on_press, on_release=on_release, daemon=True).start()
except Exception as e:
    print(f"[warn] hotkey listener off: {e}", flush=True)


# ---- consent modal (tkinter, topmost) ----
def show_modal(code, manifest, active_fn):
    import tkinter as tk
    root = tk.Tk(); root.title("CUA Tutor — consent required")
    root.attributes("topmost", True); root.attributes("alpha", 0.97)
    w, h = 560, 380
    root.geometry(f"{w}x{h}+{root.winfo_screenwidth()//2 - w//2}+120")
    root.configure(bg="#101418")
    tk.Label(root, text="🎓 CUA Tutor wants to assist you", fg="#e6edf3",
             bg="#101418", font=("Segoe UI", 16, "bold")).pack(pady=(18, 4))
    tk.Label(root, text=f"Confirm code:\n{code}", fg="#f0b429", bg="#101418",
             font=("Consolas", 30, "bold"), justify="center").pack(pady=8)
    m = manifest or {}
    tk.Label(root, text=("Allowed apps: " + ", ".join(m.get("apps", [])) +
                         f"\nAllowed actions: {', '.join(m.get('actions', []))}"
                         f"\nBudget: {m.get('budget')} actions · "
                         f"{m.get('duration_min')} minutes"
                         "\nType this code on the web page to allow · "
                         "Ctrl+Alt+F12 stops everything"),
             fg="#8b949e", bg="#101418", justify="left",
             font=("Segoe UI", 10)).pack(pady=4)
    st = tk.Label(root, text="waiting for your confirmation…", fg="#58a6ff",
                  bg="#101418")
    st.pack(pady=6)
    tk.Button(root, text="🛑 DENY / STOP", bg="#da3633", fg="white",
              command=lambda: die("user denied")).pack(pady=8)

    def tick():
        if STATE["kill"]:
            root.destroy(); return
        if active_fn():
            root.destroy(); return
        root.after(700, tick)
    root.after(700, tick)
    root.mainloop()


def screenshot_b64():
    img = pyautogui.screenshot()
    img.thumbnail((1280, 1280))
    buf = io.BytesIO()
    img.convert("RGB").save(buf, "JPEG", quality=55)
    return base64.b64encode(buf.getvalue()).decode()


def execute(act):
    try:
        kind = act.get("kind")
        if kind == "click":
            sw, sh = pyautogui.size()
            pyautogui.click(int(act["x"]) / 1000 * sw, int(act["y"]) / 1000 * sh)
        elif kind == "type":
            pyautogui.typewrite(str(act.get("text", "")), interval=0.03)
        elif kind == "key":
            pyautogui.hotkey(*str(act.get("combo", "")).split("+"))
        elif kind == "scroll":
            pyautogui.scroll(int(act.get("amount", 3)))
        elif kind == "wait":
            time.sleep(min(float(act.get("sec", 1)), 5))
        return True, "ok"
    except Exception as e:
        return False, f"{type(e).__name__}: {e}"


print(f"[agent] session {A.sid} @ {BASE} (Ctrl+Alt+F12 or mouse-to-top-left = STOP)", flush=True)
modal_open = {"v": False}
state_now = {"v": "waiting"}


def poll_state():
    r = requests.get(f"{URL}/poll", params=K, timeout=40)
    r.raise_for_status()
    return r.json()


while not STATE["kill"]:
    try:
        d = poll_state()
    except Exception as e:
        print(f"[warn] poll: {e}", flush=True); time.sleep(3); continue
    state_now["v"] = d.get("state")
    if state_now["v"] == "await_consent":
        if not modal_open["v"]:
            def runner(c=d["code"], m=d["manifest"]):
                modal_open["v"] = True
                show_modal(c, m, lambda: STATE["kill"] or state_now["v"] not in ("await_consent",))
                modal_open["v"] = False
            threading.Thread(target=runner, daemon=True).start()
        continue
    if state_now["v"] in ("killed", "done", "expired"):
        print(f"[end] {state_now['v']}", flush=True); break
    act = d.get("action")
    if act and act.get("id"):
        ok, det = execute(act)
        requests.post(f"{URL}/result", params=K, json={
            "action_id": act["id"], "ok": ok, "detail": det,
            "screenshot": screenshot_b64()}, timeout=30)
print("bye")