"""
Build K-step GUI Transition data from AndroidControl dataset for GUI-Shift training.
Run this first to construct the training dataset.
"""
import json
import os
import random
from collections import defaultdict
from datasets import load_dataset

def parse_action(action_str):
    """Parse AndroidControl action string to GUI-Shift action format."""
    if not action_str.startswith("!FUNCTIONCALL"):
        return None
    json_str = action_str.replace("!FUNCTIONCALL", "")
    try:
        action = json.loads(json_str)
        name = action["name"]
        params = action.get("parameters", {})
        if name == "click":
            return {"action_type": "click", "x": params["x"], "y": params["y"]}
        elif name == "long_press":
            return {"action_type": "long_press", "x": params["x"], "y": params["y"]}
        elif name == "scroll":
            return {"action_type": "scroll", "direction": params["direction"]}
        elif name == "open_app":
            return {"action_type": "open_app", "app_name": params["app_name"]}
        elif name == "input_text":
            return {"action_type": "input_text", "text": params["text"]}
        elif name == "navigate_back":
            return {"action_type": "navigate_back"}
        elif name == "navigate_home":
            return {"action_type": "navigate_home"}
        elif name == "wait":
            return {"action_type": "wait"}
        elif name == "status":
            return None
        else:
            return None
    except Exception:
        return None

def build_k_step_transitions(dataset_split, k=1, max_samples=2000, output_dir="./data"):
    os.makedirs(output_dir, exist_ok=True)
    episodes = defaultdict(list)
    for i, item in enumerate(dataset_split):
        json_data = item["json"]
        episode_id = json_data["episode_id"]
        episodes[episode_id].append({
            "index": i,
            "step_id": json_data["step_id"],
            "messages": json_data["messages"],
            "png": item["png"],
        })
    for ep_id in episodes:
        episodes[ep_id].sort(key=lambda x: x["step_id"])

    transitions = []
    for ep_id, steps in episodes.items():
        for i in range(len(steps) - k):
            state_t = steps[i]
            state_tk = steps[i + k]
            messages = state_t["messages"]
            action_str = None
            for msg in messages:
                if msg["role"] == "assistant":
                    action_str = msg["content"]
                    break
            if not action_str:
                continue
            action = parse_action(action_str)
            if action is None:
                continue

            img_t = state_t["png"]
            img_tk = state_tk["png"]
            img_t_path = os.path.join(output_dir, f"ep{ep_id}_step{state_t['step_id']}.png")
            img_tk_path = os.path.join(output_dir, f"ep{ep_id}_step{state_tk['step_id']}.png")
            img_t.save(img_t_path)
            img_tk.save(img_tk_path)

            problem = (
                "You are given two GUI screenshots. The first screenshot shows the current state, "
                "and the second screenshot shows the state after executing some actions. "
                "Predict the first action that transitions from the first screenshot to the second screenshot."
            )
            transitions.append({
                "image": [img_t_path, img_tk_path],
                "conversations": [
                    {"from": "human", "value": f"<image>\n<image>\n{problem}"},
                    {"from": "gpt", "value": json.dumps(action)}
                ],
                "problem": problem,
                "solution": json.dumps(action),
            })

    random.shuffle(transitions)
    transitions = transitions[:max_samples]
    output_path = os.path.join(output_dir, f"ui_transition_training_{max_samples}_k_{k}_no_reasoning.jsonl")
    with open(output_path, "w") as f:
        for item in transitions:
            f.write(json.dumps(item) + "\n")
    print(f"Saved {len(transitions)} K-step transitions to {output_path}")
    return output_path

if __name__ == "__main__":
    print("Loading AndroidControl dataset...")
    ds = load_dataset("ckg/AndroidControlParsedWithImages-20k", split="train", streaming=True)
    items = []
    for i, item in enumerate(ds):
        items.append(item)
        if i >= 10000:
            break
    print(f"Processing {len(items)} items...")
    output = build_k_step_transitions(items, k=1, max_samples=2000, output_dir="./data")
    print(f"Done: {output}")