"""Evaluate a trained agent on MiniWoB++ tasks."""

import json
import argparse
import re

from agent_system.environments.env_package.miniwob.miniwob_env import MiniWoBEnv
from agent_system.environments.prompts.miniwob import SYSTEM_PROMPT, build_messages


EVAL_TASKS = [
    "miniwob/click-button-v1",
    "miniwob/click-dialog-v1",
    "miniwob/click-link-v1",
    "miniwob/click-checkboxes-v1",
    "miniwob/enter-text-v1",
    "miniwob/enter-password-v1",
    "miniwob/login-user-v1",
    "miniwob/navigate-tree-v1",
    "miniwob/search-engine-v1",
    "miniwob/social-media-v1",
]

NUM_EPISODES_PER_TASK = 50


def parse_action_from_response(response: str) -> str:
    """从模型 response 中提取 <action>...</action> 标签内的动作。"""
    m = re.search(r'<action>(.*?)</action>', response, re.DOTALL)
    if m:
        return m.group(1).strip()
    # fallback: 尝试直接匹配 click/type/press 模式
    m = re.search(r'(click|type|press)\(.*?\)', response)
    if m:
        return m.group(0)
    # 无法解析时返回 noop
    return "click(0, 0)"


def evaluate(model, task_name: str, num_episodes: int = 50) -> float:
    """在单个任务上评估 num_episodes 个 episode，返回 success rate。"""
    env = MiniWoBEnv(task_name=task_name, max_steps=10)
    successes = 0

    for ep in range(num_episodes):
        obs = env.reset()
        done = False
        history = []
        step = 0

        while not done:
            messages = build_messages(obs, history, step)
            response = model.predict(messages)
            action_str = parse_action_from_response(response)

            prev_obs = obs
            obs, reward, done, info = env.step(action_str)
            history.append({"action": action_str, "obs": prev_obs})
            step += 1

        if env.total_reward > 0:
            successes += 1

    env.close()
    return successes / num_episodes


def load_model(checkpoint_path: str):
    """从 verl-agent checkpoint 加载模型。实际实现依赖 checkpoint 格式。"""
    raise NotImplementedError(
        f"Please implement load_model() for checkpoint: {checkpoint_path}\n"
        "The model should expose a predict(messages) -> str interface."
    )


def main():
    parser = argparse.ArgumentParser(description="Evaluate trained MiniWoB++ agent")
    parser.add_argument("--checkpoint", type=str, required=True, help="Path to model checkpoint")
    parser.add_argument("--output", type=str, default="eval_results.json", help="Output JSON file")
    parser.add_argument("--tasks", nargs="+", default=None, help="Task subset to evaluate (default: all)")
    parser.add_argument("--num_episodes", type=int, default=NUM_EPISODES_PER_TASK)
    args = parser.parse_args()

    model = load_model(args.checkpoint)
    tasks = args.tasks or EVAL_TASKS

    results = {}
    for task in tasks:
        sr = evaluate(model, task, args.num_episodes)
        results[task] = {"success_rate": sr, "num_episodes": args.num_episodes}
        print(f"{task}: SR = {sr:.2%}")

    avg_sr = sum(r["success_rate"] for r in results.values()) / len(results)
    results["average"] = {"success_rate": avg_sr}
    print(f"\nAverage SR: {avg_sr:.2%}")

    with open(args.output, "w") as f:
        json.dump(results, f, indent=2)
    print(f"Results saved to {args.output}")


if __name__ == "__main__":
    main()