"""AGUVIS stage 2 (smolagents/aguvis-stage-2): computer-use steps with normalized coordinates, from
AndroidControl, Mind2Web, GUIAct, GUI-Odyssey, MiniWoB, AMEX and AitW.

Each source's stream is read from row `skip` on, `n / 8` usable steps per source, then shuffled with
`split_seed`. The shuffle depends on n, so two pools of different sizes hold different rows in
different orders: compare runs on test_steps() (rows 6,000+ of every source, past every training pool),
not on the first rows of a pool.
"""

import itertools, random, re

from . import load_cache, save_cache

AGUVIS = ["android_control", "mind2web-l2", "guiact-web-single", "guiact-web-multi-l2", "gui-odyssey-l2",
          "miniwob-l2", "amex-l2", "aitw-l2"]
ACTIONS = ["click", "type", "scroll", "swipe", "press", "open", "back", "home", "wait", "answer", "select"]
KEYS = ["enter", "tab", "escape", "space", "delete", "up", "down", "left", "right", "ctrl", "other"]  # all single tokens
KEY_ALIAS = {"backspace": "delete", "esc": "escape", "return": "enter"}
SDIR = ["up", "down", "left", "right"]
CALL = re.compile(r"<code>(.*?)</code>", re.S)
FN = re.compile(r"(\w+)\((.*?)\)\s*$", re.S)


def parse_action(assistant):
    """The assistant's <code> call as (action, {x, y, x2, y2, key, sdir, samt}), or None."""
    m = CALL.search(assistant or "")
    if not m:
        return None
    body = m.group(1).strip().splitlines()
    call = FN.search(body[-1].strip()) if body else None
    if not call:
        return None
    fn, args = call.group(1), call.group(2)

    def num(k):
        m = re.search(rf"{k}\s*=\s*([-\d.]+)", args)
        return float(m.group(1)) if m else None

    def pair(k):
        m = re.search(rf"{k}\s*=\s*\[(.*?)\]", args)
        return [float(v) for v in re.findall(r"[-\d.]+", m.group(1))] if m else None

    out = {}
    if fn in ("click", "select_option"):
        out.update(x=num("x"), y=num("y"))
        act = "click" if fn == "click" else "select"
    elif fn == "swipe":
        a, b = pair("from_coord"), pair("to_coord")
        if not a or not b:
            return None
        out.update(x=a[0], y=a[1], x2=b[0], y2=b[1])
        act = "swipe"
    elif fn == "scroll":
        d = re.search(r"direction\s*=\s*['\"](\w+)['\"]", args)
        amt = num("amount")
        out.update(sdir=d.group(1) if d else None, samt=min(9, int((amt or 5) / 2)))
        act = "scroll"
    elif fn == "press":
        k = re.search(r"keys?\s*=\s*['\"]([^'\"]+)['\"]", args) or re.search(r"['\"]([^'\"]+)['\"]", args)
        key = k.group(1).lower().split("+")[-1] if k else "other"
        key = KEY_ALIAS.get(key, key)
        out["key"] = key if key in KEYS else "other"
        act = "press"
    else:
        act = {"type": "type", "open_app": "open", "navigate_back": "back", "navigate_home": "home",
               "wait": "wait", "final_answer": "answer", "terminate": "answer"}.get(fn)
        if act is None:
            return None
    if act in ("click", "select", "swipe") and not all(0 <= (out.get(k) or -1) <= 1 for k in ("x", "y")):
        return None
    return act, out


def steps(n, skip=0, split_seed=0, cache=True):
    """n AGUVIS steps as dict(image, user, action, args, source)."""
    name = f"aguvis_{n}_{skip}_{split_seed}"
    if cache and (got := load_cache(name)) is not None:
        return got
    from datasets import load_dataset
    out, per = [], max(1, n // len(AGUVIS))
    for cfg in AGUVIS:
        ds = load_dataset("smolagents/aguvis-stage-2", cfg, split="train", streaming=True)
        got = 0
        for r in itertools.islice(ds, skip, skip + per * 6):
            if got >= per:  # enough from this source: stop streaming (the rows kept are the same)
                break
            if len(r["images"]) != 1:
                continue
            t = r["texts"][0]
            a = parse_action(t.get("assistant"))
            if not a:
                continue
            out.append(dict(image=r["images"][0], user=(t.get("user") or "").strip(), action=a[0], args=a[1],
                            source=cfg))
            got += 1
        print(f"  {cfg}: {got}", flush=True)
    random.Random(split_seed).shuffle(out)
    out = out[:n]
    if cache:
        save_cache(name, out)
    return out


def test_steps(n=1600):
    """The clean held-out test set: rows 6,000+ of every source's stream, past every training pool
    (a pool of N reads at most 6 N / 8 rows per source)."""
    return steps(1600, skip=6000, split_seed=1)[: n or None]