"""Wave-UI (agentsea/wave-ui): UI elements with instructions and boxes, used as clicks in the agent
layout and as element boxes. The ScreenSpot source and every image that also appears in ScreenSpot
are removed, so ScreenSpot stays a clean test set."""

import hashlib

from . import load_cache, save_cache

# sources whose "instruction" is element-type metadata; their element name reads like an instruction
NAME_SOURCES = ("webui", "roboflow", "mind2web_test_domain", "mind2web_test_task", "mind2web_test_website")


def screenspot_hashes():
    from datasets import load_dataset
    key = lambda im: hashlib.md5(im.convert("L").resize((64, 64)).tobytes()).hexdigest()
    return key, {key(im) for im in load_dataset("bevaya/ScreenSpot", split="test")["image"]}


def clean(split, n, seed, key, banned):
    from datasets import load_dataset
    ds = load_dataset("agentsea/wave-ui", split=split).shuffle(seed=seed)
    ds = ds.filter(lambda src: src != "screenspot", input_columns=["source"])
    out = []
    for r in ds:
        if len(out) >= n:
            break
        x1, y1, x2, y2 = r["bbox"]
        W, H = r["image"].size
        if not r["instruction"] or x2 <= x1 or y2 <= y1 or x2 > W * 1.02 or y2 > H * 1.02:
            continue
        if key(r["image"]) in banned:
            continue
        if r["source"] in NAME_SOURCES:
            r = dict(r, instruction=r["name"])
        if not r["instruction"]:
            continue
        out.append(r)
    return out


def clicks(n, skip=0, cache=True):
    """Wave-UI elements as click steps in the agent layout (keeps the pointing sharp)."""
    name = f"waveui_clicks_{n}_{skip}"
    if cache and (got := load_cache(name)) is not None:
        return got
    key, banned = screenspot_hashes()
    out = []
    for r in clean("train", skip + n, 0, key, banned)[skip:]:
        W, H = r["image"].size
        x1, y1, x2, y2 = r["bbox"]
        out.append(dict(image=r["image"], user=f"Instruction: {r['instruction']}\n\nPrevious actions:\nNone",
                        action="click", source="wave-ui",
                        args=dict(x=(x1 + x2) / 2 / W, y=(y1 + y2) / 2 / H, box=(x2 - x1) / W, boy=(y2 - y1) / H)))
    if cache:
        save_cache(name, out)
    return out


def boxes(n, skip=2000, cache=True):
    """Wave-UI element boxes (pixels). Rows [skip, skip + n) are disjoint from the clicks [0, n_click)."""
    name = f"waveui_boxes_{n}_{skip}"
    if cache and (got := load_cache(name)) is not None:
        return got
    key, banned = screenspot_hashes()
    out = [dict(image=r["image"], instruction=r["instruction"], box=list(r["bbox"]), source=r["source"])
           for r in clean("train", skip + n, 0, key, banned)[skip:]]
    if cache:
        save_cache(name, out)
    return out


def val_boxes(n):
    key, banned = screenspot_hashes()
    return [dict(image=r["image"], instruction=r["instruction"], box=list(r["bbox"]), source=r["source"])
            for r in clean("validation", n, 1, key, banned)]