#!/usr/bin/env python3
"""Vision-driven Android control loop demo.
screencap -> vision model picks tap coords (JSON) -> adb input tap -> verify screen changed.
Usage: python3 android_control_demo.py <model> <goal> [host]
"""
import base64, json, os, re, sys, subprocess, urllib.request, time

ADB = "/opt/homebrew/bin/adb"
SERIAL = os.environ.get("ANDROID_SERIAL", "Android.local:5555")
HOST = sys.argv[3] if len(sys.argv) > 3 else "http://10.30.20.186:11434"
MODEL = sys.argv[1]
GOAL = sys.argv[2]
TMP = os.path.expanduser("~/vision_bench/ctl")

def sh(*args, **kw):
    return subprocess.run(list(args), capture_output=True, text=True, timeout=kw.get("timeout", 30))

def screencap(path):
    r = subprocess.run([ADB, "-s", SERIAL, "exec-out", "screencap", "-p"],
                       capture_output=True, timeout=30)  # binary, no text decode
    if r.returncode != 0 or not r.stdout:
        return False
    with open(path, "wb") as f:
        f.write(r.stdout)
    return True

def ask_vision(model, img_path, goal, history):
    b64 = base64.b64encode(open(img_path, "rb").read()).decode()
    hist_txt = ("\n".join(f"step {i+1}: {h}" for i, h in enumerate(history[-6:])) if history
                else "(no steps taken yet)")
    prompt = (f'You control this Android phone. The screenshot is 1024x768 pixels; your tap '
              f'coordinates must be in that same pixel space. GOAL: "{goal}". '
              f'Steps already taken:\n{hist_txt}\n'
              'Look at the screen and decide the SINGLE next step toward the goal. '
              'Output ONLY one valid JSON object, no commentary: '
              '{"action": "tap", "coordinate": [<int x>, <int y>], "reason": "<one line>"}. '
              'Use action "back" or "home" when needed, or "done" if the goal is achieved.')
    payload = {"model": model, "stream": False, "keep_alive": "10m",
               "options": {"temperature": 0, "num_predict": 200, "num_ctx": int(os.environ.get("NUM_CTX", "2048"))},
               "messages": [{"role": "user", "content": prompt, "images": [b64]}]}
    req = urllib.request.Request(HOST + "/api/chat", data=json.dumps(payload).encode(),
                                 headers={"Content-Type": "application/json"})
    opener = urllib.request.build_opener(urllib.request.ProxyHandler({}))
    t0 = time.time()
    with opener.open(req, timeout=180) as r:
        d = json.loads(r.read())
    content = d["message"].get("content", "")
    if not content:
        content = d["message"].get("thinking", "")
    m = None
    objs = re.findall(r"\{[^{}]*\}", content)
    for o in objs:
        try:
            j = json.loads(o)
        except Exception:
            continue
        if not isinstance(j, dict):
            continue
        # prefer explicit coordinate
        coord = j.get("coordinate") or j.get("center") or j.get("point")
        x, y = j.get("x"), j.get("y")
        if isinstance(coord, list) and len(coord) >= 2:
            x, y = int(coord[0]), int(coord[1])
        elif isinstance(x, list) and isinstance(y, list) and x and y:
            x, y = int(x[0]), int(y[0])  # bbox-ish -> top-left of each axis
        elif isinstance(x, (int, float)) and isinstance(y, (int, float)):
            x, y = int(x), int(y)
        else:
            continue
        m = {"action": j.get("action", "tap"), "x": x, "y": y,
             "reason": j.get("reason", "")}
        if m["action"] in ("click", "left_click"):
            m["action"] = "tap"
        break
    return m, content, round(time.time() - t0, 1)

def main():
    os.makedirs(TMP, exist_ok=True)
    sh(ADB, "-s", SERIAL, "connect", SERIAL, timeout=10)
    history = []
    for step in range(1, 9):
        img = f"{TMP}/step{step}.png"
        if not screencap(img):
            print("screencap failed"); return
        action, raw, wall = ask_vision(MODEL, img, GOAL, history)
        print(f"[step {step}] {wall}s -> {json.dumps(action)}")
        if not action:
            print("  RAW:", raw[:200]); return
        if action.get("action") == "done":
            print("GOAL REACHED per model"); return
        if action.get("action") == "swipe":
            x, y = action.get("x"), action.get("y")
            if x is not None and y is not None:
                sh(ADB, "-s", SERIAL, "shell", "input", "swipe", str(x), str(y), str(x), str(max(0, y-400)), "300")
                history.append(f"swiped up at ({x},{y}) — {action.get('reason','')}")
            else:
                history.append("swipe skipped (no coords)")
        elif action.get("action") == "back":
            sh(ADB, "-s", SERIAL, "shell", "input", "keyevent", "4")
            history.append("pressed back")
        elif action.get("action") == "home":
            sh(ADB, "-s", SERIAL, "shell", "input", "keyevent", "3")
            history.append("pressed home")
        else:
            x, y = action.get("x"), action.get("y")
            if x is None or y is None:
                print("no coords, RAW:", raw[:200]); return
            sh(ADB, "-s", SERIAL, "shell", "input", "tap", str(x), str(y))
            history.append(f"tapped ({x},{y}) — {action.get('reason','')}")
        time.sleep(2.5)
    print("MAX STEPS REACHED")

if __name__ == "__main__":
    main()