Blame for frozen/trace2task/visual_nanojev_data.py

"""Offline contracts, not a verified adapter for the raw AgentNet schema.

Callers must establish screenshot timing and source coordinate semantics first.
No source code, screenshot, or desktop action is executed here.
"""

from __future__ import annotations

import ast
import copy
import hashlib
import json
import math
import random
from collections.abc import Mapping

from .actions import SPECIAL_KEYS, ActionCall

# Source recorder spellings are retained, not silently mapped across OSes.
SOURCE_KEYS = tuple(
    sorted(
        SPECIAL_KEYS
        | set(
            "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
            " !\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~"
        )
        | {
            "cmd",
            "command",
            "win",
            "winleft",
            "winright",
            "esc",
            "caps_lock",
            "capslock",
            "num_lock",
            "numlock",
            "scroll_lock",
            "scrolllock",
            "print_screen",
            "printscreen",
            "menu",
            "pause",
            "media_volume_up",
            "media_volume_down",
            "media_volume_mute",
            "ctrl_l",
            "ctrl_r",
            "alt_l",
            "alt_r",
            "shift_l",
            "shift_r",
            "cmd_l",
            "cmd_r",
            "altleft",
            "altright",
            "ctrlleft",
            "ctrlright",
            "shiftleft",
            "shiftright",
            "ż",
            "ś",
            "ł",
            *(f"f{number}" for number in range(1, 25)),
        }
    )
)

ACTIONS = (
    "click",
    "double_click",
    "scroll",
    "press_key",
    "hotkey",
    "type_text",
    "drag",
    "wait",
    "done",
    "end_step",
    "move_to",
    "drag_to",
    "wheel",
    "horizontal_wheel",
    "triple_click",
    "key_down",
    "key_up",
)


def validate_action(payload):
    if payload == {"done": True} and type(payload.get("done")) is bool:
        return {"done": True}
    if not isinstance(payload, dict) or set(payload) != {"skill", "args"}:
        raise ValueError("Expected one exact action object")
    skill, args = payload["skill"], payload["args"]
    if skill not in ACTIONS or skill == "done":
        raise ValueError("Unsupported action")
    if skill == "end_step":
        if args != {}:
            raise ValueError("end_step takes no arguments")
        return {"skill": skill, "args": {}}
    if skill in {"move_to", "drag_to", "triple_click"}:
        required = {"x", "y"} | ({"button"} if skill != "move_to" else set())
        if not isinstance(args, dict) or set(args) != required:
            raise ValueError("Unexpected pointer arguments")
        point = ActionCall("click", args).args
        if skill == "move_to":
            point.pop("button")
        return {"skill": skill, "args": point}
    if skill in {"wheel", "horizontal_wheel"}:
        if not isinstance(args, dict) or set(args) != {"wheel_units"}:
            raise ValueError("Wheel requires an amount")
        units = args["wheel_units"]
        if type(units) is not int or units == 0:
            raise ValueError("Invalid wheel units")
        return {"skill": skill, "args": {"wheel_units": units}}
    if skill == "type_text":
        if not isinstance(args, dict) or set(args) != {"text"} or not isinstance(args["text"], str):
            raise ValueError("Text action requires a string")
        # Source text is not clipped, stripped, or stripped of newlines. Token
        # limits are checked separately by the processor before training.
        return {"skill": skill, "args": dict(args)}
    if skill in {"press_key", "hotkey", "key_down", "key_up"}:
        expected = "keys" if skill == "hotkey" else "key"
        if not isinstance(args, dict) or set(args) != {expected}:
            raise ValueError("Unexpected keyboard arguments")
        keys = args["keys"] if skill == "hotkey" else [args["key"]]
        if not isinstance(keys, list) or not keys or any(key not in SOURCE_KEYS for key in keys):
            raise ValueError("Unsupported source keyboard key")
        if skill == "hotkey" and (not 2 <= len(keys) <= 4 or len(set(keys)) != len(keys)):
            raise ValueError("Unsupported hotkey arity or duplicate keys")
        return {"skill": skill, "args": copy.deepcopy(args)}
    if skill == "scroll":
        if not isinstance(args, dict) or set(args) != {"x", "y", "wheel_units"}:
            raise ValueError("Scroll requires location and explicit wheel units")
        position = ActionCall("click", {k: args[k] for k in ("x", "y")}).args
        units = args["wheel_units"]
        if type(units) is not int or units == 0 or not -100 <= units <= 100:
            raise ValueError("Invalid wheel units")
        return {
            "skill": skill,
            "args": {"x": position["x"], "y": position["y"], "wheel_units": units},
        }
    return ActionCall.from_payload(payload).to_payload()


def validate_step(payload):
    """A source step is one or more ordered actions under ONE observation."""
    if isinstance(payload, dict) and set(payload) == {"actions"}:
        actions = payload["actions"]
        if not isinstance(actions, list) or not actions:
            raise ValueError("A step requires a nonempty action sequence")
        clean = [validate_action(action) for action in actions]
        if any(action.get("skill") == "end_step" for action in clean):
            raise ValueError("end_step is a decoder boundary, not a source action")
        if any(action.get("done") for action in clean[:-1]):
            raise ValueError("No action may follow termination")
        return {"actions": clean}
    return validate_action(payload)


def step_actions(payload):
    value = validate_step(payload)
    return value.get("actions", [value])


def validate_history_action(payload):
    if isinstance(payload, dict) and set(payload) == {"source_code"}:
        if not isinstance(payload["source_code"], str) or not payload["source_code"]:
            raise ValueError("Missing recorded history code")
        return dict(payload)
    return validate_step(payload)


def encode_coordinate(value, grid=32):
    if type(value) not in (int, float) or not math.isfinite(value) or not 0 <= value <= 1:
        raise ValueError("Expected a finite normalized coordinate")
    if type(grid) is not int or grid < 1:
        raise ValueError("Invalid grid")
    cell = min(grid - 1, math.floor(grid * value))
    return cell, grid * value - cell


def pixel_coordinate(value, dimension):
    encode_coordinate(value)
    if type(dimension) is not int or dimension < 1:
        raise ValueError("Invalid image dimension")
    return min(dimension - 1, math.floor(value * dimension))


def parse_literal_action(code):
    """Parse our compact interchange language, NOT arbitrary pyautogui/AgentNet code.

    Keyword-only calls make units and fields explicit. Raw source adaptation must
    separately verify aliases, omitted defaults, timing and coordinate conventions.
    """
    if not isinstance(code, str) or len(code) > 10000:
        raise ValueError("Invalid action expression")
    try:
        node = ast.parse(code, mode="eval").body
        if (
            not isinstance(node, ast.Call)
            or not isinstance(node.func, ast.Name)
            or node.func.id not in ACTIONS
            or node.args
        ):
            raise ValueError("Only a single keyword-only whitelisted call is allowed")
        args = {}
        for keyword in node.keywords:
            if keyword.arg is None or keyword.arg in args:
                raise ValueError("Expanded or duplicate arguments are forbidden")
            args[keyword.arg] = ast.literal_eval(keyword.value)
        if node.func.id == "done":
            if args:
                raise ValueError("done takes no arguments")
            return {"done": True}
        return validate_action({"skill": node.func.id, "args": args})
    except (SyntaxError, TypeError, RecursionError) as exc:
        raise ValueError("Invalid literal action") from exc


def serialize_action(payload, arm):
    action = validate_action(payload)
    if arm == "G-json":
        return json.dumps(action, ensure_ascii=False, separators=(",", ":"), allow_nan=False)
    if arm == "G-compact":
        if action.get("done"):
            return "done()"
        return (
            action["skill"] + "(" + ",".join(f"{k}={v!r}" for k, v in action["args"].items()) + ")"
        )
    raise ValueError("Unknown text baseline")


def build_record(
    *,
    record_id,
    trajectory_id,
    duplicate_group,
    step_index,
    task,
    screenshot,
    history,
    target,
    audit_only,
    synthetic=False,
):
    """Build from explicit reviewed fields, never merge untrusted source dictionaries."""
    for value in (record_id, trajectory_id, duplicate_group, task):
        if not isinstance(value, str) or not value.strip():
            raise ValueError("Identity, duplicate grouping and original task are required")
    if type(step_index) is not int or step_index < 0 or type(synthetic) is not bool:
        raise ValueError("Invalid step/synthetic marker")
    fields = {"asset_id", "sha256", "width", "height", "timing"}
    if not isinstance(screenshot, dict) or set(screenshot) != fields:
        raise ValueError("Unexpected screenshot fields")
    if screenshot["timing"] != "before_target_action":
        raise ValueError("Unverified screenshot timing")
    if not isinstance(screenshot["asset_id"], str) or not screenshot["asset_id"]:
        raise ValueError("Missing image identity")
    for key in ("width", "height"):
        pixel_coordinate(0, screenshot[key])
    digest = screenshot["sha256"]
    if not synthetic and (
        not isinstance(digest, str)
        or len(digest) != 64
        or any(c not in "0123456789abcdef" for c in digest)
    ):
        raise ValueError("Real images require verified SHA256")
    clean_history = []
    previous = -1
    for item in history:
        if set(item) != {"step_index", "action", "executed"}:
            raise ValueError("Unexpected history fields")
        index = item["step_index"]
        if (
            type(index) is not int
            or not previous < index < step_index
            or item["executed"] is not True
        ):
            raise ValueError("History must contain ordered previously executed actions")
        action = validate_history_action(item["action"])
        if "source_code" not in action and any(a.get("done") for a in step_actions(action)):
            raise ValueError("History cannot continue after done")
        clean_history.append({"step_index": index, "action": action, "executed": True})
        previous = index
    target = validate_step(target)
    if any(a.get("done") for a in step_actions(target)) and not (
        audit_only.get("terminal_state_observed") is True
        or audit_only.get("terminal_label_basis")
        == "source_demonstration_stop_not_verified_success"
    ):
        raise ValueError("done supervision requires an observed terminal state")
    return {
        "record_id": record_id,
        "trajectory_id": trajectory_id,
        "duplicate_group": duplicate_group,
        "step_index": step_index,
        "synthetic": synthetic,
        "input": {
            "task": task,
            "screenshot": copy.deepcopy(screenshot),
            "history": clean_history[-4:],
        },
        "target": target,
        "audit_only": copy.deepcopy(audit_only),
    }


def policy_input(record):
    """Only this allowlisted object may be passed to a future backbone processor."""
    value = record["input"]
    if set(value) != {"task", "screenshot", "history"}:
        raise ValueError("Unexpected policy input fields")
    return copy.deepcopy(value)


def grouped_splits(records, salt="visual-nanojev-v1"):
    """Connected trajectory/template/exact-image groups; no near-duplicate detector.

    duplicate_group must be supplied by a separate reviewed grouping procedure.
    Four deterministic partitions are development/train/calibration/internal_test.
    """
    parent = list(range(len(records)))

    def root(i):
        while parent[i] != i:
            parent[i] = parent[parent[i]]
            i = parent[i]
        return i

    seen = {}
    ids = set()
    for i, record in enumerate(records):
        if record["record_id"] in ids:
            raise ValueError("Duplicate record identity")
        ids.add(record["record_id"])
        keys = [("trajectory", record["trajectory_id"]), ("template", record["duplicate_group"])]
        digest = record["input"]["screenshot"]["sha256"]
        if digest:
            keys.append(("image", digest))
        for key in keys:
            if key in seen:
                parent[root(i)] = root(seen[key])
            seen[key] = i
    groups = {}
    for i, record in enumerate(records):
        groups.setdefault(root(i), []).append(record["record_id"])
    assignments = {}
    for members in groups.values():
        identity = json.dumps([salt, sorted(members)], separators=(",", ":"))
        bucket = int(hashlib.sha256(identity.encode()).hexdigest(), 16) % 100
        split = (
            "train"
            if bucket < 80
            else "development"
            if bucket < 90
            else "calibration"
            if bucket < 95
            else "internal_test"
        )
        assignments.update(dict.fromkeys(members, split))
    return assignments


def trajectory_splits(records, seed=42, validation_fraction=0.05):
    """Assign complete trajectories to deterministic train/validation splits.

    Trajectory identities are sorted before fixed-seed sampling, so assignments do
    not depend on record input order.  The validation trajectory count is
    ``ceil(n * validation_fraction)`` and is clipped to leave at least one
    trajectory in each split when ``n >= 2``.  Zero or one trajectory is assigned
    entirely to train because a disjoint validation split is then impossible.

    This intentionally performs no template, image, or similarity grouping.
    """
    if type(seed) is not int:
        raise ValueError("Seed must be an integer")
    if (
        type(validation_fraction) not in (int, float)
        or not math.isfinite(validation_fraction)
        or not 0 < validation_fraction < 1
    ):
        raise ValueError("Validation fraction must be finite and strictly between zero and one")

    record_trajectories = {}
    trajectory_ids = set()
    for record in records:
        if not isinstance(record, Mapping):
            raise TypeError("Each record must be a mapping")
        record_id = record.get("record_id")
        trajectory_id = record.get("trajectory_id")
        if not isinstance(record_id, str) or not record_id.strip():
            raise ValueError("Each record requires a non-empty string record_id")
        if not isinstance(trajectory_id, str) or not trajectory_id.strip():
            raise ValueError("Each record requires a non-empty string trajectory_id")
        if record_id in record_trajectories:
            if record_trajectories[record_id] != trajectory_id:
                raise ValueError("Conflicting duplicate record identity")
            raise ValueError("Duplicate record identity")
        record_trajectories[record_id] = trajectory_id
        trajectory_ids.add(trajectory_id)

    ordered_trajectories = sorted(trajectory_ids)
    validation_count = 0
    if len(ordered_trajectories) >= 2:
        validation_count = min(
            len(ordered_trajectories) - 1,
            max(1, math.ceil(len(ordered_trajectories) * validation_fraction)),
        )
    validation = set(random.Random(seed).sample(ordered_trajectories, validation_count))
    return {
        record_id: "validation" if trajectory_id in validation else "train"
        for record_id, trajectory_id in record_trajectories.items()
    }