scripts/run_agentlab_miniwob_study.py

#!/usr/bin/env python3
"""Run a real AgentLab MiniWoB study for a local LoRA adapter checkpoint.

This script is the "official-comparison" path that complements
`estimate_miniwob_adapter_score.py`:

- prepares a LoRA adapter archive/directory for inference
- launches an AgentLab MiniWoB study
- collects the study summary and `study_id`
- writes a leaderboard-style `miniwob.json`
- writes a draft `README.md` scaffold for BrowserGym leaderboard submission

It is still your responsibility to verify the metadata fields before submitting
the generated files upstream.
"""

from __future__ import annotations

import argparse
import csv
import json
import logging
import os
import re
import statistics
import sys
from collections import Counter
from datetime import datetime, timezone
from fnmatch import fnmatch
from pathlib import Path
from typing import Any, Dict, List, Mapping, Optional, Sequence
from urllib.parse import quote

ROOT = Path(__file__).resolve().parents[1]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))


SLICE_PRESETS: Dict[str, Dict[str, Any]] = {
    "full": {
        "description": "Full MiniWoB benchmark.",
        "task_include": [],
        "task_exclude": [],
        "seeds_per_task": None,
        "task_limit": None,
        "episode_limit": None,
    },
    "smoke": {
        "description": "Very small real MiniWoB slice for quick sanity checks.",
        "task_include": [
            "ascending-numbers",
            "choose-list",
            "click-link",
            "click-option",
            "click-color",
            "click-dialog*",
        ],
        "task_exclude": [],
        "seeds_per_task": 1,
        "task_limit": 6,
        "episode_limit": 6,
    },
    "starter": {
        "description": "Cheap real MiniWoB slice across a broader spread of task families.",
        "task_include": [
            "ascending-numbers",
            "bisect-angle",
            "book-flight*",
            "choose-date*",
            "choose-list",
            "click-button-sequence",
            "click-checkboxes*",
            "click-color",
            "click-dialog*",
            "click-link",
            "click-option",
            "click-pie",
        ],
        "task_exclude": [],
        "seeds_per_task": 1,
        "task_limit": 12,
        "episode_limit": 12,
    },
    "core": {
        "description": "Moderate real MiniWoB slice with two seeds per selected task family.",
        "task_include": [
            "ascending-numbers",
            "bisect-angle",
            "book-flight*",
            "choose-date*",
            "choose-list",
            "click-button-sequence",
            "click-checkboxes*",
            "click-color",
            "click-dialog*",
            "click-link",
            "click-option",
            "click-pie",
        ],
        "task_exclude": [],
        "seeds_per_task": 2,
        "task_limit": 12,
        "episode_limit": 24,
    },
    "clicks": {
        "description": "Click-heavy real MiniWoB slice.",
        "task_include": [
            "ascending-numbers",
            "circle-center",
            "click-button-sequence",
            "click-color",
            "click-dialog*",
            "click-link",
            "click-menu*",
            "click-pie",
        ],
        "task_exclude": [],
        "seeds_per_task": 1,
        "task_limit": 8,
        "episode_limit": 8,
    },
    "forms": {
        "description": "Form and selection-oriented real MiniWoB slice.",
        "task_include": [
            "book-flight*",
            "buy-ticket",
            "choose-date*",
            "choose-list",
            "click-checkboxes*",
            "click-option",
        ],
        "task_exclude": [],
        "seeds_per_task": 1,
        "task_limit": 8,
        "episode_limit": 8,
    },
}

LEADERBOARD_SPACE_ID = "ServiceNow/browsergym-leaderboard"


def _safe_float(value: Any, default: float = 0.0) -> float:
    try:
        if value is None:
            return default
        if isinstance(value, str) and not value.strip():
            return default
        numeric = float(value)
        if numeric != numeric:
            return default
        return numeric
    except Exception:
        return default


def _safe_int(value: Any, default: int = 0) -> int:
    return int(round(_safe_float(value, float(default))))


def _safe_mean(values: Sequence[float]) -> float:
    clean = [float(value) for value in values if value is not None]
    return statistics.fmean(clean) if clean else 0.0


def _append_query_param(url: str, key: str, value: str) -> str:
    if not url or not value:
        return url
    separator = "&" if "?" in url else "?"
    return f"{url}{separator}{key}={quote(value)}"


def build_trackio_context(
    *,
    project: str = "",
    run_name: str = "",
    space_id: str = "",
    dashboard_url: str = "",
    enabled: bool = False,
    available: bool = False,
) -> Dict[str, Any]:
    resolved_project = (project or "").strip()
    resolved_run_name = (run_name or "").strip()
    resolved_space_id = (
        (space_id or "").strip()
        or os.getenv("TRACKIO_SPACE_ID", "").strip()
        or os.getenv("BROWSER_ENV_TRACKIO_SPACE_ID", "").strip()
    )
    resolved_dashboard_url = (dashboard_url or os.getenv("TRACKIO_DASHBOARD_URL", "")).strip()
    if not resolved_dashboard_url and resolved_space_id:
        resolved_dashboard_url = f"https://{resolved_space_id.replace('/', '-').lower()}.hf.space"
    project_dashboard_url = (
        _append_query_param(resolved_dashboard_url, "project", resolved_project)
        if resolved_dashboard_url and resolved_project
        else resolved_dashboard_url
    )
    space_page_url = f"https://huggingface.co/spaces/{resolved_space_id}" if resolved_space_id else ""
    return {
        "enabled": bool(enabled),
        "available": bool(available),
        "configured": bool(enabled or resolved_project or resolved_space_id or resolved_dashboard_url),
        "project": resolved_project,
        "run_name": resolved_run_name,
        "space_id": resolved_space_id,
        "dashboard_url": resolved_dashboard_url,
        "project_dashboard_url": project_dashboard_url,
        "space_page_url": space_page_url,
        "link_url": project_dashboard_url or resolved_dashboard_url or space_page_url,
    }


def find_column(
    table: Any,
    *,
    preferred: Sequence[str] = (),
    suffixes: Sequence[str] = (),
    contains: Sequence[str] = (),
) -> Optional[str]:
    columns = [str(column) for column in getattr(table, "columns", [])]
    for candidate in preferred:
        if candidate in columns:
            return candidate
    for column in columns:
        if any(column.endswith(suffix) for suffix in suffixes):
            return column
    for column in columns:
        if any(token in column for token in contains):
            return column
    return None


def classify_error_message(message: str) -> str:
    text = str(message or "").strip()
    if not text:
        return ""
    policy_match = re.search(r"(policy_error|translation_error):([A-Za-z0-9_]+)", text)
    if policy_match:
        return f"{policy_match.group(1)}:{policy_match.group(2)}"
    error_match = re.search(r"([A-Za-z_][A-Za-z0-9_]*(?:Error|Exception|Failure))", text)
    if error_match:
        return error_match.group(1)
    if ":" in text:
        return text.split(":", 1)[0].strip()[:80]
    return text[:80]


def truncate_text(value: Any, limit: int = 160) -> str:
    text = str(value or "").strip()
    if len(text) <= limit:
        return text
    return text[: limit - 3].rstrip() + "..."


def sanitize_metric_slug(value: str) -> str:
    return re.sub(r"[^a-z0-9]+", "-", str(value or "").strip().lower()).strip("-") or "task"


def minimize_row(row: Mapping[str, Any]) -> Dict[str, Any]:
    return {
        "agent_name": str(row.get("agent_name", "")),
        "score": _safe_float(row.get("score")),
        "std_err": _safe_float(row.get("std_err")),
        "comments": truncate_text(row.get("comments", ""), limit=200),
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--adapter-source", required=True, help="Path to final_adapter directory or zip archive")
    parser.add_argument(
        "--base-model-override",
        default="",
        help="Optional explicit base model override. If blank, prefer manifest base_model_id, then adapter_config.",
    )
    parser.add_argument("--agent-name", default="final-adapter-agentlab")
    parser.add_argument(
        "--submission-agent-name",
        default="",
        help="Optional leaderboard display name. Defaults to --agent-name.",
    )
    parser.add_argument("--model-name", default="", help="Override the model name written to the README scaffold.")
    parser.add_argument("--work-dir", default="artifacts/agentlab_miniwob_study")
    parser.add_argument("--study-suffix", default="miniwob-official")
    parser.add_argument(
        "--comment",
        default="MiniWoB AgentLab study for browser_env LoRA adapter.",
        help="Comment stored in AgentLab reproducibility info.",
    )
    parser.add_argument("--n-jobs", type=int, default=1, help="Parallel AgentLab jobs. Keep 1 for single-GPU local adapter evals.")
    parser.add_argument(
        "--parallel-backend",
        choices=["sequential", "joblib", "ray"],
        default="sequential",
    )
    parser.add_argument("--n-relaunch", type=int, default=1)
    parser.add_argument("--strict-reproducibility", action="store_true")
    parser.add_argument("--max-new-tokens", type=int, default=72)
    parser.add_argument("--max-elements", type=int, default=24)
    parser.add_argument("--temperature", type=float, default=0.0)
    parser.add_argument("--device-map", default="auto")
    parser.add_argument("--torch-dtype", default="auto", choices=["auto", "float16", "bfloat16"])
    parser.add_argument("--load-in-4bit", action="store_true", default=True)
    parser.add_argument("--no-4bit", dest="load_in_4bit", action="store_false")
    parser.add_argument(
        "--benchmark-specific",
        choices=["Yes", "No", "Unknown"],
        default="Unknown",
    )
    parser.add_argument(
        "--benchmark-tuned",
        choices=["Yes", "No", "Unknown"],
        default="Unknown",
    )
    parser.add_argument(
        "--reproducible",
        choices=["Yes", "No", "Unknown"],
        default="Unknown",
    )
    parser.add_argument(
        "--original-or-reproduced",
        choices=["Original", "Reproduced"],
        default="Original",
    )
    parser.add_argument("--paper-link", default="")
    parser.add_argument("--code-repository", default="")
    parser.add_argument("--license-name", default="")
    parser.add_argument("--dataset-used", default="TODO")
    parser.add_argument("--training-steps", default="TODO")
    parser.add_argument("--hardware-used", default="TODO")
    parser.add_argument("--training-time", default="TODO")
    parser.add_argument("--additional-notes", default="")
    parser.add_argument("--trackio", action="store_true")
    parser.add_argument("--trackio-project", default="browser-rl-openenv-agentlab")
    parser.add_argument("--trackio-run-name", default="")
    parser.add_argument(
        "--trackio-space-id",
        default="",
        help="Optional Hugging Face Space id for hosted Trackio dashboards, e.g. org/space.",
    )
    parser.add_argument(
        "--slice-preset",
        choices=sorted(SLICE_PRESETS),
        default="full",
        help="Run a smaller real MiniWoB slice before paying for the full benchmark sweep.",
    )
    parser.add_argument(
        "--task-include",
        action="append",
        default=[],
        help="Optional glob matched against task family names like `click-color` or full task names like `miniwob.click-color_12`.",
    )
    parser.add_argument(
        "--task-exclude",
        action="append",
        default=[],
        help="Optional glob matched against task family names or full task names to exclude from the run.",
    )
    parser.add_argument(
        "--seeds-per-task",
        type=int,
        default=None,
        help="Optional per-task-family seed cap after slice filtering.",
    )
    parser.add_argument(
        "--task-limit",
        type=int,
        default=None,
        help="Optional cap on the number of task families kept after filtering.",
    )
    parser.add_argument(
        "--episode-limit",
        type=int,
        default=None,
        help="Optional cap on the total task-seed episodes kept after filtering.",
    )
    parser.add_argument(
        "--list-slice",
        action="store_true",
        help="Print the selected slice and exit without launching AgentLab.",
    )
    parser.add_argument("--output", default="")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

    from scripts.estimate_miniwob_adapter_score import prepare_adapter

    work_dir = Path(args.work_dir).expanduser().resolve()
    work_dir.mkdir(parents=True, exist_ok=True)

    prepared_dir = prepare_adapter(
        Path(args.adapter_source).expanduser(),
        work_dir / "prepared_adapter",
        base_model_override="",
    )
    adapter_metadata = load_adapter_metadata(prepared_dir)
    chosen_base_model = choose_base_model(
        explicit_override=args.base_model_override,
        manifest=adapter_metadata.get("adapter_manifest"),
        adapter_config=adapter_metadata["adapter_config"],
    )
    if chosen_base_model:
        from scripts.estimate_miniwob_adapter_score import apply_base_model_override

        apply_base_model_override(prepared_dir, chosen_base_model)
        adapter_metadata["effective_base_model"] = chosen_base_model
    else:
        adapter_metadata["effective_base_model"] = str(
            adapter_metadata["adapter_config"].get("base_model_name_or_path", "")
        )

    os.environ.setdefault("AGENTLAB_EXP_ROOT", str(work_dir / "agentlab_results"))
    verify_lora_runtime(prepared_dir)

    tracker_bundle = maybe_init_trackio(args, adapter_metadata=adapter_metadata, prepared_dir=prepared_dir)
    study, result_df, summary_df, error_report, benchmark_meta = run_agentlab_study(
        args,
        prepared_dir=prepared_dir,
        effective_base_model=str(adapter_metadata["effective_base_model"]),
    )
    summary_record = extract_summary_record(summary_df)
    result_meta = summarize_result_df(result_df)
    analysis = build_eval_analysis(
        summary_df=summary_df,
        result_df=result_df,
        summary_record=summary_record,
        benchmark_meta=benchmark_meta,
    )
    leaderboard_row = build_leaderboard_row(
        args=args,
        study=study,
        summary_record=summary_record,
        result_meta=result_meta,
        benchmark_meta=benchmark_meta,
    )
    leaderboard_comparison = compare_with_live_leaderboard(
        float(leaderboard_row["score"]),
        benchmark_meta=benchmark_meta,
    )

    submission_agent_name = args.submission_agent_name or args.agent_name
    submission_slug = slugify(submission_agent_name)
    submission_dir = work_dir / "browsergym_submission" / "results" / submission_slug
    submission_dir.mkdir(parents=True, exist_ok=True)

    readme_text = build_submission_readme(
        args=args,
        submission_agent_name=submission_agent_name,
        summary_record=summary_record,
        result_meta=result_meta,
        study=study,
        adapter_metadata=adapter_metadata,
        benchmark_meta=benchmark_meta,
    )
    (submission_dir / "README.md").write_text(readme_text, encoding="utf-8")
    (submission_dir / "miniwob.json").write_text(json.dumps([leaderboard_row], indent=2) + "\n", encoding="utf-8")

    analysis_files = write_analysis_artifacts(work_dir / "analysis", analysis)

    output_path = Path(args.output or (work_dir / "study_summary.json")).expanduser().resolve()
    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_payload = {
        "created_at": datetime.now(timezone.utc).isoformat(),
        "study_id": str(study.uuid),
        "study_dir": str(getattr(study, "dir", "")),
        "agent_name": args.agent_name,
        "submission_agent_name": submission_agent_name,
        "leaderboard_row": leaderboard_row,
        "result_meta": result_meta,
        "summary_record": summary_record,
        "benchmark_meta": benchmark_meta,
        "adapter_metadata": adapter_metadata,
        "analysis": analysis,
        "analysis_files": analysis_files,
        "leaderboard_comparison": leaderboard_comparison,
        "trackio": tracker_bundle["context"],
        "error_report_path": str(Path(getattr(study, "dir", "")) / "error_report.md"),
        "submission_dir": str(submission_dir),
        "notes": build_output_notes(benchmark_meta),
    }
    output_path.write_text(json.dumps(output_payload, indent=2) + "\n", encoding="utf-8")

    print(json.dumps(output_payload["leaderboard_row"], indent=2))
    print(f"\nStudy ID: {study.uuid}")
    print(f"Study dir: {getattr(study, 'dir', '')}")
    print(f"Submission scaffold: {submission_dir}")
    print(f"Summary JSON: {output_path}")
    print(f"Task summary CSV: {analysis_files['task_summary_csv']}")
    if leaderboard_comparison.get("available"):
        print(
            "Projected leaderboard rank: "
            f"{leaderboard_comparison.get('projected_rank')} / {leaderboard_comparison.get('field_size')}"
        )
    if tracker_bundle["context"].get("link_url"):
        print(f"Trackio dashboard: {tracker_bundle['context']['link_url']}")

    tracker = tracker_bundle["client"]
    if tracker is not None:
        try:
            metrics: Dict[str, float] = {
                "agentlab/miniwob_score": float(leaderboard_row["score"]),
                "agentlab/miniwob_std_err": float(leaderboard_row["std_err"]),
                "agentlab/avg_reward": float(summary_record.get("avg_reward", 0.0)),
                "agentlab/avg_steps": float(summary_record.get("avg_steps", 0.0)),
                "agentlab/task_count": float(result_meta["task_count"]),
                "agentlab/episode_count": float(result_meta["episode_count"]),
                "agentlab/error_count": float(result_meta["error_count"]),
                "analysis/error_rate": float(analysis["aggregate"]["error_rate"]),
                "analysis/avg_parse_failures_per_episode": float(
                    analysis["aggregate"]["avg_parse_failures_per_episode"]
                ),
                "analysis/avg_invalid_actions_per_episode": float(
                    analysis["aggregate"]["avg_invalid_actions_per_episode"]
                ),
                "analysis/avg_repeated_actions_per_episode": float(
                    analysis["aggregate"]["avg_repeated_actions_per_episode"]
                ),
                "analysis/avg_episode_policy_latency_ms": float(
                    analysis["aggregate"]["avg_episode_policy_latency_ms"]
                ),
            }
            if leaderboard_comparison.get("available"):
                metrics["leaderboard/projected_rank"] = float(
                    leaderboard_comparison.get("projected_rank", 0)
                )
                metrics["leaderboard/field_size"] = float(leaderboard_comparison.get("field_size", 0))
            for task_row in analysis.get("per_task", []):
                slug = sanitize_metric_slug(str(task_row.get("task_family", "")))
                metrics[f"tasks/{slug}/score_pct"] = float(task_row.get("score_pct", 0.0))
                metrics[f"tasks/{slug}/avg_steps"] = float(task_row.get("avg_steps", 0.0))
                metrics[f"tasks/{slug}/error_rate"] = float(task_row.get("error_rate", 0.0))
                metrics[f"tasks/{slug}/avg_parse_failures"] = float(
                    task_row.get("avg_parse_failures", 0.0)
                )
                metrics[f"tasks/{slug}/avg_invalid_actions"] = float(
                    task_row.get("avg_invalid_actions", 0.0)
                )
                metrics[f"tasks/{slug}/avg_episode_policy_latency_ms"] = float(
                    task_row.get("avg_episode_policy_latency_ms", 0.0)
                )
            tracker.log(metrics)
        finally:
            tracker.finish()


def load_adapter_metadata(adapter_dir: Path) -> Dict[str, Any]:
    config_path = adapter_dir / "adapter_config.json"
    if not config_path.exists():
        raise FileNotFoundError(f"Missing adapter_config.json in {adapter_dir}")
    manifest_path = adapter_dir / "browser_env_adapter_manifest.json"
    payload: Dict[str, Any] = {
        "adapter_dir": str(adapter_dir),
        "adapter_config": json.loads(config_path.read_text(encoding="utf-8")),
    }
    if manifest_path.exists():
        payload["adapter_manifest"] = json.loads(manifest_path.read_text(encoding="utf-8"))
    return payload


def choose_base_model(
    *,
    explicit_override: str,
    manifest: Optional[Mapping[str, Any]],
    adapter_config: Mapping[str, Any],
) -> str:
    if explicit_override:
        return explicit_override
    manifest_model = ""
    if isinstance(manifest, Mapping):
        manifest_model = str(manifest.get("base_model_id", "") or "").strip()
    if manifest_model:
        return manifest_model
    return str(adapter_config.get("base_model_name_or_path", "") or "").strip()


def run_agentlab_study(
    args: argparse.Namespace,
    *,
    prepared_dir: Path,
    effective_base_model: str,
):
    from scripts.estimate_miniwob_adapter_score import choose_torch_dtype

    try:
        from bgym import DEFAULT_BENCHMARKS
        from agentlab.experiments.study import make_study
    except Exception as exc:
        raise RuntimeError(
            "AgentLab is required for a real MiniWoB study. Install it with `pip install agentlab` before running this script."
        ) from exc
    try:
        from browser_env.agentlab_miniwob_adapter import MiniWoBAdapterAgentArgs
    except ImportError:
        from agentlab_miniwob_adapter import MiniWoBAdapterAgentArgs

    dtype = choose_torch_dtype(args.torch_dtype)
    logging_level = logging.INFO
    agent_args = MiniWoBAdapterAgentArgs(
        adapter_dir=str(prepared_dir),
        base_model_override=effective_base_model,
        agent_name=args.agent_name,
        max_new_tokens=args.max_new_tokens,
        max_elements=args.max_elements,
        temperature=args.temperature,
        load_in_4bit=args.load_in_4bit,
        device_map=args.device_map,
        torch_dtype=dtype,
    )
    benchmark = DEFAULT_BENCHMARKS["miniwob"]()
    benchmark, benchmark_meta = apply_miniwob_slice(benchmark, args)
    if args.list_slice:
        print_selected_slice(benchmark_meta)
        raise SystemExit(0)
    study = make_study(
        benchmark=benchmark,
        agent_args=[agent_args],
        comment=args.comment,
        suffix=resolve_study_suffix(args.study_suffix, benchmark_meta),
        logging_level=logging_level,
        logging_level_stdout=logging_level,
    )
    study.run(
        n_jobs=args.n_jobs,
        parallel_backend=args.parallel_backend,
        strict_reproducibility=args.strict_reproducibility,
        n_relaunch=args.n_relaunch,
        exp_root=os.environ.get("AGENTLAB_EXP_ROOT"),
    )
    result_df, summary_df, error_report = study.get_results()
    return study, result_df, summary_df, error_report, benchmark_meta


def verify_lora_runtime(prepared_dir: Path) -> None:
    if not (prepared_dir / "adapter_config.json").exists():
        return
    try:
        from browser_env.slm_policy import assert_hf_lora_runtime_available
    except ImportError:
        from slm_policy import assert_hf_lora_runtime_available

    assert_hf_lora_runtime_available()


def extract_summary_record(summary_df: Any) -> Dict[str, Any]:
    if summary_df is None or getattr(summary_df, "empty", False):
        raise RuntimeError("AgentLab returned an empty summary_df.")
    reset = summary_df.reset_index()
    candidates = reset.to_dict(orient="records")
    if not candidates:
        raise RuntimeError("Could not extract any summary rows from summary_df.")

    for record in candidates:
        for key, value in record.items():
            if str(value) == "[ALL TASKS]" or str(value).lower() == "miniwob":
                return sanitize_record(record)
    if len(candidates) == 1:
        return sanitize_record(candidates[0])

    best = max(candidates, key=lambda row: float(row.get("avg_reward", float("-inf")) or float("-inf")))
    return sanitize_record(best)


def summarize_result_df(result_df: Any) -> Dict[str, Any]:
    if result_df is None or getattr(result_df, "empty", False):
        return {"episode_count": 0, "task_count": 0, "error_count": 0}
    reset = result_df.reset_index()
    columns = set(str(col) for col in reset.columns)
    task_col = "env.task_name" if "env.task_name" in columns else next(
        (col for col in reset.columns if str(col).endswith("task_name")),
        None,
    )
    err_col = "err_msg" if "err_msg" in columns else next(
        (col for col in reset.columns if str(col).endswith("err_msg")),
        None,
    )
    task_count = int(reset[task_col].nunique()) if task_col is not None else 0
    error_count = int(reset[err_col].notnull().sum()) if err_col is not None else 0
    return {
        "episode_count": int(len(reset)),
        "task_count": task_count,
        "error_count": error_count,
    }


def sanitize_record(record: Mapping[str, Any]) -> Dict[str, Any]:
    clean: Dict[str, Any] = {}
    for key, value in record.items():
        if hasattr(value, "item"):
            try:
                value = value.item()
            except Exception:
                pass
        clean[str(key)] = value
    return clean


def summarize_task_entry(row: Mapping[str, Any]) -> Dict[str, Any]:
    return {
        "rank": _safe_int(row.get("rank")),
        "task_family": str(row.get("task_family", "")),
        "score_pct": round(_safe_float(row.get("score_pct")), 4),
        "avg_reward": round(_safe_float(row.get("avg_reward")), 6),
        "std_err_pct": round(_safe_float(row.get("std_err_pct")), 4),
        "avg_steps": round(_safe_float(row.get("avg_steps")), 4),
        "episode_count": _safe_int(row.get("episode_count")),
        "error_count": _safe_int(row.get("error_count")),
        "error_rate": round(_safe_float(row.get("error_rate")), 4),
        "avg_parse_failures": round(_safe_float(row.get("avg_parse_failures")), 4),
        "avg_invalid_actions": round(_safe_float(row.get("avg_invalid_actions")), 4),
        "avg_repeated_actions": round(_safe_float(row.get("avg_repeated_actions")), 4),
        "avg_episode_policy_latency_ms": round(_safe_float(row.get("avg_episode_policy_latency_ms")), 4),
        "priority_score": round(_safe_float(row.get("priority_score")), 4),
        "attention_reasons": list(row.get("attention_reasons", [])),
    }


def top_task_entries(
    per_task: Sequence[Mapping[str, Any]],
    *,
    key: str,
    reverse: bool,
    limit: int = 10,
) -> List[Dict[str, Any]]:
    ranked = sorted(
        per_task,
        key=lambda row: (_safe_float(row.get(key)), str(row.get("task_family", ""))),
        reverse=reverse,
    )
    return [summarize_task_entry(row) for row in ranked[:limit]]


def build_eval_analysis(
    *,
    summary_df: Any,
    result_df: Any,
    summary_record: Mapping[str, Any],
    benchmark_meta: Mapping[str, Any],
) -> Dict[str, Any]:
    summary_by_task: Dict[str, Dict[str, Any]] = {}
    if summary_df is not None and not getattr(summary_df, "empty", False):
        summary_reset = summary_df.reset_index()
        summary_task_col = find_column(
            summary_reset,
            preferred=["env.task_name", "task_name", "task", "index"],
            suffixes=["task_name"],
            contains=["task"],
        )
        if summary_task_col is None and len(summary_reset.columns) > 0:
            summary_task_col = str(summary_reset.columns[0])
        avg_reward_col = find_column(summary_reset, preferred=["avg_reward"], suffixes=["avg_reward"])
        avg_steps_col = find_column(summary_reset, preferred=["avg_steps"], suffixes=["avg_steps"])
        std_err_col = find_column(summary_reset, preferred=["std_err"], suffixes=["std_err"])
        for record in summary_reset.to_dict(orient="records"):
            label = str(record.get(summary_task_col, "") if summary_task_col is not None else "").strip()
            if not label:
                continue
            if label == "[ALL TASKS]" or label.lower() == "miniwob":
                continue
            task_family = miniwob_task_family(label)
            summary_by_task[task_family] = {
                "task_family": task_family,
                "summary_label": label,
                "avg_reward": _safe_float(record.get(avg_reward_col)) if avg_reward_col is not None else 0.0,
                "avg_steps": _safe_float(record.get(avg_steps_col)) if avg_steps_col is not None else 0.0,
                "std_err_pct": (
                    _safe_float(record.get(std_err_col)) * 100.0 if std_err_col is not None else 0.0
                ),
            }

    diag_by_task: Dict[str, Dict[str, Any]] = {}
    error_categories: Counter[str] = Counter()
    error_messages: Counter[str] = Counter()
    action_error_categories: Counter[str] = Counter()
    if result_df is not None and not getattr(result_df, "empty", False):
        result_reset = result_df.reset_index()
        task_col = find_column(
            result_reset,
            preferred=["env.task_name", "task_name"],
            suffixes=["task_name"],
            contains=["task"],
        )
        err_col = find_column(result_reset, preferred=["err_msg"], suffixes=["err_msg"])
        think_col = find_column(result_reset, preferred=["think"], suffixes=["think"], contains=["think"])
        policy_latency_col = find_column(
            result_reset,
            preferred=["cum_policy_latency_ms", "stats.cum_policy_latency_ms"],
            suffixes=["cum_policy_latency_ms"],
        ) or find_column(
            result_reset,
            preferred=["policy_latency_ms", "stats.policy_latency_ms"],
            suffixes=["policy_latency_ms"],
        )
        parse_fail_col = find_column(
            result_reset,
            preferred=["cum_parse_failures", "stats.cum_parse_failures"],
            suffixes=["cum_parse_failures"],
        )
        invalid_col = find_column(
            result_reset,
            preferred=["cum_invalid_actions", "stats.cum_invalid_actions"],
            suffixes=["cum_invalid_actions"],
        )
        repeated_col = find_column(
            result_reset,
            preferred=["cum_repeated_actions", "stats.cum_repeated_actions"],
            suffixes=["cum_repeated_actions"],
        )
        observed_col = find_column(
            result_reset,
            preferred=["max_observed_elements", "stats.max_observed_elements"],
            suffixes=["max_observed_elements"],
        )
        step_col = find_column(
            result_reset,
            preferred=["max_step_index", "stats.max_step_index"],
            suffixes=["max_step_index"],
        )

        for record in result_reset.to_dict(orient="records"):
            task_name = str(record.get(task_col, "") if task_col is not None else "").strip()
            if not task_name:
                continue
            task_family = miniwob_task_family(task_name)
            bucket = diag_by_task.setdefault(
                task_family,
                {
                    "task_family": task_family,
                    "episode_count": 0,
                    "error_count": 0,
                    "total_parse_failures": 0.0,
                    "total_invalid_actions": 0.0,
                    "total_repeated_actions": 0.0,
                    "total_policy_latency_ms": 0.0,
                    "total_max_observed_elements": 0.0,
                    "total_final_step_index": 0.0,
                    "error_categories": Counter(),
                    "action_error_categories": Counter(),
                },
            )
            bucket["episode_count"] += 1
            bucket["total_parse_failures"] += _safe_float(record.get(parse_fail_col)) if parse_fail_col else 0.0
            bucket["total_invalid_actions"] += _safe_float(record.get(invalid_col)) if invalid_col else 0.0
            bucket["total_repeated_actions"] += _safe_float(record.get(repeated_col)) if repeated_col else 0.0
            bucket["total_policy_latency_ms"] += _safe_float(record.get(policy_latency_col)) if policy_latency_col else 0.0
            bucket["total_max_observed_elements"] += _safe_float(record.get(observed_col)) if observed_col else 0.0
            bucket["total_final_step_index"] += _safe_float(record.get(step_col)) if step_col else 0.0

            err_msg = truncate_text(record.get(err_col, "")) if err_col is not None else ""
            if err_msg:
                category = classify_error_message(err_msg)
                bucket["error_count"] += 1
                if category:
                    bucket["error_categories"][category] += 1
                    error_categories[category] += 1
                error_messages[err_msg] += 1

            think_msg = str(record.get(think_col, "") if think_col is not None else "").strip()
            if think_msg and (
                "policy_error:" in think_msg
                or "translation_error:" in think_msg
                or "Error" in think_msg
                or "Exception" in think_msg
            ):
                action_category = classify_error_message(think_msg)
                if action_category:
                    bucket["action_error_categories"][action_category] += 1
                    action_error_categories[action_category] += 1

    per_task: List[Dict[str, Any]] = []
    for task_family in sorted(set(summary_by_task) | set(diag_by_task)):
        summary_bits = summary_by_task.get(task_family, {})
        diag_bits = diag_by_task.get(task_family, {})
        episode_count = _safe_int(diag_bits.get("episode_count"))
        error_count = _safe_int(diag_bits.get("error_count"))
        score_pct = round(_safe_float(summary_bits.get("avg_reward")) * 100.0, 4)
        error_rate = (error_count / episode_count) if episode_count else 0.0
        avg_parse_failures = (
            _safe_float(diag_bits.get("total_parse_failures")) / episode_count if episode_count else 0.0
        )
        avg_invalid_actions = (
            _safe_float(diag_bits.get("total_invalid_actions")) / episode_count if episode_count else 0.0
        )
        avg_repeated_actions = (
            _safe_float(diag_bits.get("total_repeated_actions")) / episode_count if episode_count else 0.0
        )
        avg_episode_policy_latency_ms = (
            _safe_float(diag_bits.get("total_policy_latency_ms")) / episode_count if episode_count else 0.0
        )
        avg_max_observed_elements = (
            _safe_float(diag_bits.get("total_max_observed_elements")) / episode_count if episode_count else 0.0
        )
        avg_final_step_index = (
            _safe_float(diag_bits.get("total_final_step_index")) / episode_count if episode_count else 0.0
        )
        attention_reasons: List[str] = []
        if score_pct < 50.0:
            attention_reasons.append("low_score")
        if error_count > 0:
            attention_reasons.append("episode_errors")
        if avg_parse_failures > 0:
            attention_reasons.append("parse_failures")
        if avg_invalid_actions > 0:
            attention_reasons.append("invalid_actions")
        if avg_repeated_actions > 0:
            attention_reasons.append("repeated_actions")
        priority_score = (
            (100.0 - score_pct)
            + (error_rate * 35.0)
            + (avg_parse_failures * 12.0)
            + (avg_invalid_actions * 6.0)
            + (avg_repeated_actions * 2.5)
        )
        top_error_category = (
            diag_bits["error_categories"].most_common(1)[0][0]
            if diag_bits.get("error_categories")
            else ""
        )
        top_action_error_category = (
            diag_bits["action_error_categories"].most_common(1)[0][0]
            if diag_bits.get("action_error_categories")
            else ""
        )
        per_task.append(
            {
                "task_family": task_family,
                "summary_label": summary_bits.get("summary_label", task_family),
                "avg_reward": round(_safe_float(summary_bits.get("avg_reward")), 6),
                "score_pct": score_pct,
                "std_err_pct": round(_safe_float(summary_bits.get("std_err_pct")), 4),
                "avg_steps": round(_safe_float(summary_bits.get("avg_steps")), 4),
                "episode_count": episode_count,
                "error_count": error_count,
                "error_rate": round(error_rate, 4),
                "total_parse_failures": round(_safe_float(diag_bits.get("total_parse_failures")), 4),
                "avg_parse_failures": round(avg_parse_failures, 4),
                "total_invalid_actions": round(_safe_float(diag_bits.get("total_invalid_actions")), 4),
                "avg_invalid_actions": round(avg_invalid_actions, 4),
                "total_repeated_actions": round(_safe_float(diag_bits.get("total_repeated_actions")), 4),
                "avg_repeated_actions": round(avg_repeated_actions, 4),
                "avg_episode_policy_latency_ms": round(avg_episode_policy_latency_ms, 4),
                "avg_max_observed_elements": round(avg_max_observed_elements, 4),
                "avg_final_step_index": round(avg_final_step_index, 4),
                "top_error_category": top_error_category,
                "top_action_error_category": top_action_error_category,
                "attention_reasons": attention_reasons,
                "priority_score": round(priority_score, 4),
            }
        )

    ranked_by_score = sorted(
        per_task,
        key=lambda row: (_safe_float(row.get("score_pct")), str(row.get("task_family", ""))),
        reverse=True,
    )
    for index, row in enumerate(ranked_by_score, start=1):
        row["rank"] = index

    total_episode_count = sum(_safe_int(row.get("episode_count")) for row in per_task)
    total_error_count = sum(_safe_int(row.get("error_count")) for row in per_task)
    total_parse_failures = sum(_safe_float(row.get("total_parse_failures")) for row in per_task)
    total_invalid_actions = sum(_safe_float(row.get("total_invalid_actions")) for row in per_task)
    total_repeated_actions = sum(_safe_float(row.get("total_repeated_actions")) for row in per_task)
    aggregate = {
        "benchmark_label": str(benchmark_meta.get("benchmark_label", "MiniWoB")),
        "score": round(_safe_float(summary_record.get("avg_reward")) * 100.0, 4),
        "std_err": round(_safe_float(summary_record.get("std_err")) * 100.0, 4),
        "avg_reward": round(_safe_float(summary_record.get("avg_reward")), 6),
        "avg_steps": round(_safe_float(summary_record.get("avg_steps")), 6),
        "task_family_count": len(per_task),
        "episode_count": total_episode_count,
        "error_count": total_error_count,
        "error_rate": round((total_error_count / total_episode_count) if total_episode_count else 0.0, 4),
        "task_families_with_errors": sum(_safe_int(row.get("error_count")) > 0 for row in per_task),
        "total_parse_failures": round(total_parse_failures, 4),
        "avg_parse_failures_per_episode": round(
            (total_parse_failures / total_episode_count) if total_episode_count else 0.0,
            4,
        ),
        "total_invalid_actions": round(total_invalid_actions, 4),
        "avg_invalid_actions_per_episode": round(
            (total_invalid_actions / total_episode_count) if total_episode_count else 0.0,
            4,
        ),
        "total_repeated_actions": round(total_repeated_actions, 4),
        "avg_repeated_actions_per_episode": round(
            (total_repeated_actions / total_episode_count) if total_episode_count else 0.0,
            4,
        ),
        "avg_episode_policy_latency_ms": round(
            _safe_mean([_safe_float(row.get("avg_episode_policy_latency_ms")) for row in per_task]),
            4,
        ),
        "avg_max_observed_elements": round(
            _safe_mean([_safe_float(row.get("avg_max_observed_elements")) for row in per_task]),
            4,
        ),
        "avg_final_step_index": round(
            _safe_mean([_safe_float(row.get("avg_final_step_index")) for row in per_task]),
            4,
        ),
    }
    return {
        "aggregate": aggregate,
        "per_task": ranked_by_score,
        "rankings": {
            "best_by_score": top_task_entries(ranked_by_score, key="score_pct", reverse=True),
            "worst_by_score": top_task_entries(ranked_by_score, key="score_pct", reverse=False),
            "highest_error_rate": top_task_entries(ranked_by_score, key="error_rate", reverse=True),
            "highest_parse_failures": top_task_entries(ranked_by_score, key="avg_parse_failures", reverse=True),
            "highest_invalid_actions": top_task_entries(ranked_by_score, key="avg_invalid_actions", reverse=True),
            "slowest_tasks": top_task_entries(
                ranked_by_score,
                key="avg_episode_policy_latency_ms",
                reverse=True,
            ),
            "improvement_priority": top_task_entries(ranked_by_score, key="priority_score", reverse=True),
        },
        "error_breakdown": {
            "episode_error_categories": [
                {"category": category, "count": count}
                for category, count in error_categories.most_common()
            ],
            "episode_error_messages": [
                {"message": message, "count": count}
                for message, count in error_messages.most_common(20)
            ],
            "action_error_categories": [
                {"category": category, "count": count}
                for category, count in action_error_categories.most_common()
            ],
        },
    }


def write_analysis_artifacts(analysis_dir: Path, analysis: Mapping[str, Any]) -> Dict[str, str]:
    analysis_dir.mkdir(parents=True, exist_ok=True)
    task_summary_json = analysis_dir / "task_summary.json"
    task_summary_csv = analysis_dir / "task_summary.csv"
    error_breakdown_json = analysis_dir / "error_breakdown.json"
    study_analysis_json = analysis_dir / "study_analysis.json"

    task_rows = list(analysis.get("per_task", []))
    task_summary_json.write_text(json.dumps(task_rows, indent=2) + "\n", encoding="utf-8")

    fieldnames = [
        "rank",
        "task_family",
        "score_pct",
        "avg_reward",
        "std_err_pct",
        "avg_steps",
        "episode_count",
        "error_count",
        "error_rate",
        "total_parse_failures",
        "avg_parse_failures",
        "total_invalid_actions",
        "avg_invalid_actions",
        "total_repeated_actions",
        "avg_repeated_actions",
        "avg_episode_policy_latency_ms",
        "avg_max_observed_elements",
        "avg_final_step_index",
        "top_error_category",
        "top_action_error_category",
        "priority_score",
        "attention_reasons",
    ]
    with task_summary_csv.open("w", encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        for row in task_rows:
            payload = dict(row)
            payload["attention_reasons"] = ",".join(str(reason) for reason in row.get("attention_reasons", []))
            writer.writerow({key: payload.get(key, "") for key in fieldnames})

    error_breakdown_json.write_text(
        json.dumps(analysis.get("error_breakdown", {}), indent=2) + "\n",
        encoding="utf-8",
    )
    study_analysis_json.write_text(json.dumps(analysis, indent=2) + "\n", encoding="utf-8")
    return {
        "task_summary_json": str(task_summary_json),
        "task_summary_csv": str(task_summary_csv),
        "error_breakdown_json": str(error_breakdown_json),
        "study_analysis_json": str(study_analysis_json),
    }


def compare_with_live_leaderboard(score: float, *, benchmark_meta: Mapping[str, Any]) -> Dict[str, Any]:
    if not benchmark_meta.get("is_full_benchmark", True):
        return {
            "available": False,
            "reason": "Partial MiniWoB slice runs are not directly comparable to the full BrowserGym MiniWoB leaderboard.",
        }
    try:
        from huggingface_hub import HfApi, hf_hub_download
    except Exception as exc:
        return {"available": False, "error": f"huggingface_hub unavailable: {exc}"}

    api = HfApi()
    try:
        repo_files = api.list_repo_files(repo_id=LEADERBOARD_SPACE_ID, repo_type="space")
        miniwob_files = [
            path for path in repo_files if path.startswith("results/") and path.endswith("/miniwob.json")
        ]
        rows: List[Dict[str, Any]] = []
        for score_file in miniwob_files:
            local_path = hf_hub_download(
                repo_id=LEADERBOARD_SPACE_ID,
                repo_type="space",
                filename=score_file,
            )
            payload = json.loads(Path(local_path).read_text(encoding="utf-8"))
            if not isinstance(payload, list):
                continue
            for item in payload:
                if item.get("benchmark") != "MiniWoB":
                    continue
                if item.get("original_or_reproduced") != "Original":
                    continue
                rows.append(dict(item))
    except Exception as exc:
        return {"available": False, "error": str(exc), "source": LEADERBOARD_SPACE_ID}

    ranked = sorted(rows, key=lambda row: (-_safe_float(row.get("score")), str(row.get("agent_name", ""))))
    projected_rank = 1 + sum(_safe_float(row.get("score")) > score for row in ranked)
    above = [row for row in ranked if _safe_float(row.get("score")) > score]
    below = [row for row in ranked if _safe_float(row.get("score")) <= score]
    payload: Dict[str, Any] = {
        "available": True,
        "source": LEADERBOARD_SPACE_ID,
        "projected_rank": projected_rank,
        "field_size": len(ranked),
        "top_entries": [minimize_row(row) for row in ranked[:10]],
    }
    if above:
        payload["nearest_above"] = minimize_row(above[-1])
    if below:
        payload["nearest_below"] = minimize_row(below[0])
    return payload


def build_leaderboard_row(
    *,
    args: argparse.Namespace,
    study: Any,
    summary_record: Mapping[str, Any],
    result_meta: Mapping[str, Any],
    benchmark_meta: Mapping[str, Any],
) -> Dict[str, Any]:
    avg_reward = float(summary_record.get("avg_reward", 0.0) or 0.0)
    std_err = float(summary_record.get("std_err", 0.0) or 0.0)
    score = avg_reward * 100.0
    score_std_err = std_err * 100.0
    benchmark_name = str(benchmark_meta.get("benchmark_label", "MiniWoB"))
    is_full_benchmark = bool(benchmark_meta.get("is_full_benchmark", True))
    protocol_label = "Yes - AgentLab MiniWoB study" if is_full_benchmark else f"Partial - {benchmark_name} slice"
    slice_comment = ""
    if not is_full_benchmark:
        slice_comment = (
            f" slice_preset={benchmark_meta.get('slice_preset')};"
            f" selected_task_families={int(benchmark_meta.get('selected_task_count', 0))};"
            f" selected_episodes={int(benchmark_meta.get('selected_env_count', 0))};"
        )
    return {
        "agent_name": args.submission_agent_name or args.agent_name,
        "study_id": str(study.uuid),
        "date_time": datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S"),
        "benchmark": benchmark_name,
        "score": round(score, 4),
        "std_err": round(score_std_err, 4),
        "benchmark_specific": args.benchmark_specific,
        "benchmark_tuned": args.benchmark_tuned,
        "followed_evaluation_protocol": protocol_label,
        "reproducible": args.reproducible,
        "comments": (
            "AgentLab MiniWoB study using browser_env.agentlab_miniwob_adapter. "
            f"benchmark={benchmark_name};"
            f"avg_reward={avg_reward:.6f}; avg_steps={float(summary_record.get('avg_steps', 0.0) or 0.0):.6f}; "
            f"task_count={int(result_meta['task_count'])}; episode_count={int(result_meta['episode_count'])}; "
            f"error_count={int(result_meta['error_count'])};{slice_comment} study_dir={getattr(study, 'dir', '')}."
        ),
        "original_or_reproduced": args.original_or_reproduced,
    }


def build_submission_readme(
    *,
    args: argparse.Namespace,
    submission_agent_name: str,
    summary_record: Mapping[str, Any],
    result_meta: Mapping[str, Any],
    study: Any,
    adapter_metadata: Mapping[str, Any],
    benchmark_meta: Mapping[str, Any],
) -> str:
    manifest = adapter_metadata.get("adapter_manifest") if isinstance(adapter_metadata, Mapping) else None
    config = adapter_metadata.get("adapter_config") if isinstance(adapter_metadata, Mapping) else {}
    model_name = args.model_name or str(
        (manifest or {}).get("base_model_id")
        or config.get("base_model_name_or_path")
        or "Unknown base model"
    )
    license_name = args.license_name or str((manifest or {}).get("license", "") or "")
    architecture_bits = [
        f"Base model: `{model_name}`",
        "Adapter: LoRA / PEFT",
    ]
    if config.get("r") is not None:
        architecture_bits.append(f"rank r={config.get('r')}")
    if config.get("lora_alpha") is not None:
        architecture_bits.append(f"alpha={config.get('lora_alpha')}")
    architecture = ", ".join(architecture_bits)
    benchmark_label = str(benchmark_meta.get("benchmark_label", "MiniWoB"))
    is_full_benchmark = bool(benchmark_meta.get("is_full_benchmark", True))

    lines = [
        f"# {submission_agent_name}",
        "",
        (
            "Auto-generated draft for BrowserGym leaderboard submission. Review and replace `TODO` fields before submitting."
            if is_full_benchmark
            else "Auto-generated draft for a partial real MiniWoB slice. This is useful for cheaper AgentLab checks, but it is not a full leaderboard submission."
        ),
        "",
        "## Required Information",
        "",
        f"- **Model Name:** {model_name}",
        f"- **Model Architecture:** {architecture}",
        (
            "- **Input/Output Format:** Input is the repo's compact `BrowserObservation` "
            "with ranked visible elements and short action history. Output is strict "
            "`BrowserAction` JSON translated into BrowserGym high-level action strings."
        ),
        "- **Training Details:**",
        f"  - Dataset used: {args.dataset_used}",
        f"  - Number of training steps: {args.training_steps}",
        f"  - Hardware used: {args.hardware_used}",
        f"  - Training time: {args.training_time}",
        "",
        "## AgentLab Study",
        "",
        f"- **study_id:** `{study.uuid}`",
        f"- **study_dir:** `{getattr(study, 'dir', '')}`",
        f"- **Benchmark:** `{benchmark_label}`",
        f"- **Tasks completed:** {result_meta['task_count']}",
        f"- **Episodes completed:** {result_meta['episode_count']}",
        f"- **Average reward:** {float(summary_record.get('avg_reward', 0.0) or 0.0):.6f}",
        f"- **Std err:** {float(summary_record.get('std_err', 0.0) or 0.0):.6f}",
        f"- **Average steps:** {float(summary_record.get('avg_steps', 0.0) or 0.0):.6f}",
        "",
    ]
    if not is_full_benchmark:
        lines.extend(
            [
                "## Slice Details",
                "",
                f"- **Slice preset:** `{benchmark_meta.get('slice_preset', 'custom')}`",
                f"- **Selected task families:** {benchmark_meta.get('selected_task_count', 0)}",
                f"- **Selected task family names:** {', '.join(benchmark_meta.get('selected_task_families', [])) or 'None'}",
                "",
            ]
        )
    lines.extend(
        [
        "## Optional Information",
        "",
        f"- **Paper Link:** {args.paper_link or 'TODO'}",
        f"- **Code Repository:** {args.code_repository or 'TODO'}",
        f"- **Additional Notes:** {args.additional_notes or 'TODO'}",
        f"- **License:** {license_name or 'TODO'}",
        "",
        "## Submission Notes",
        "",
        "- `benchmark_specific`, `benchmark_tuned`, and `reproducible` in `miniwob.json` should be reviewed manually.",
        "- This scaffold was generated from a local AgentLab run of a repo-defined custom agent.",
        (
            "- If you re-run the study, replace both this README metadata and the generated `miniwob.json` row."
            if is_full_benchmark
            else "- This run only covers a filtered MiniWoB slice. Do not submit the generated `miniwob.json` upstream as a full MiniWoB result."
        ),
        "",
        ]
    )
    return "\n".join(lines)


def maybe_init_trackio(
    args: argparse.Namespace,
    *,
    adapter_metadata: Mapping[str, Any],
    prepared_dir: Path,
) -> Any:
    run_name = args.trackio_run_name or f"agentlab-miniwob-{datetime.now(timezone.utc).strftime('%Y%m%d-%H%M%S')}"
    context = build_trackio_context(
        project=args.trackio_project if args.trackio else "",
        run_name=run_name if args.trackio else "",
        space_id=args.trackio_space_id,
        enabled=args.trackio,
        available=False,
    )
    if not args.trackio:
        return {"client": None, "context": context}

    try:
        import trackio
    except Exception as exc:
        raise RuntimeError("Trackio requested but not installed. Install `trackio` first.") from exc

    context = build_trackio_context(
        project=args.trackio_project,
        run_name=run_name,
        space_id=args.trackio_space_id,
        enabled=True,
        available=True,
    )
    trackio.init(
        project=context["project"],
        name=context["run_name"],
        space_id=context["space_id"] or None,
        config={
            "benchmark": "miniwob",
            "slice_preset": args.slice_preset,
            "task_include": list(args.task_include),
            "task_exclude": list(args.task_exclude),
            "seeds_per_task": args.seeds_per_task,
            "task_limit": args.task_limit,
            "episode_limit": args.episode_limit,
            "agent_name": args.agent_name,
            "prepared_adapter_dir": str(prepared_dir),
            "base_model": str(adapter_metadata.get("effective_base_model", "")),
            "load_in_4bit": bool(args.load_in_4bit),
            "max_new_tokens": int(args.max_new_tokens),
            "max_elements": int(args.max_elements),
            "parallel_backend": args.parallel_backend,
            "n_jobs": int(args.n_jobs),
        },
    )
    return {"client": trackio, "context": context}


def slugify(value: str) -> str:
    value = value.strip().lower()
    value = re.sub(r"[^a-z0-9]+", "-", value)
    return value.strip("-") or "agent"


def resolve_study_suffix(study_suffix: str, benchmark_meta: Mapping[str, Any]) -> str:
    if benchmark_meta.get("is_full_benchmark", True):
        return study_suffix
    slice_slug = slugify(str(benchmark_meta.get("slice_label", "slice")))
    if not study_suffix:
        return slice_slug
    if slice_slug in study_suffix:
        return study_suffix
    return f"{study_suffix}-{slice_slug}"


def miniwob_task_family(task_name: str) -> str:
    normalized = str(task_name or "").strip()
    if "." in normalized:
        normalized = normalized.split(".", 1)[1]
    maybe_family, sep, maybe_seed = normalized.rpartition("_")
    if sep and maybe_seed.isdigit():
        return maybe_family
    return normalized


def matches_task_pattern(task_name: str, task_family: str, pattern: str) -> bool:
    candidates = {task_name, task_family}
    if "." in task_name:
        candidates.add(task_name.split(".", 1)[1])
    return any(fnmatch(candidate, pattern) for candidate in candidates)


def apply_miniwob_slice(benchmark: Any, args: argparse.Namespace) -> tuple[Any, Dict[str, Any]]:
    slice_plan = resolve_slice_plan(args)
    original_env_args = list(getattr(benchmark, "env_args_list", []) or [])
    entries: List[Dict[str, Any]] = []
    for env_args in original_env_args:
        task_name = str(getattr(env_args, "task_name", "") or "")
        task_family = miniwob_task_family(task_name)
        entries.append(
            {
                "env_args": env_args,
                "task_name": task_name,
                "task_family": task_family,
            }
        )

    if not entries:
        raise RuntimeError("MiniWoB benchmark did not expose any env_args_list entries.")

    selected_entries: List[Dict[str, Any]] = []
    selected_task_families: List[str] = []
    seeds_kept_per_task: Dict[str, int] = {}
    selected_family_set = set()
    include_patterns = slice_plan["task_include"]
    exclude_patterns = slice_plan["task_exclude"]

    for entry in entries:
        task_name = entry["task_name"]
        task_family = entry["task_family"]
        if include_patterns and not any(matches_task_pattern(task_name, task_family, pattern) for pattern in include_patterns):
            continue
        if exclude_patterns and any(matches_task_pattern(task_name, task_family, pattern) for pattern in exclude_patterns):
            continue
        if slice_plan["task_limit"] is not None and task_family not in selected_family_set:
            if len(selected_family_set) >= slice_plan["task_limit"]:
                continue
        if slice_plan["seeds_per_task"] is not None and seeds_kept_per_task.get(task_family, 0) >= slice_plan["seeds_per_task"]:
            continue
        selected_entries.append(entry)
        seeds_kept_per_task[task_family] = seeds_kept_per_task.get(task_family, 0) + 1
        if task_family not in selected_family_set:
            selected_family_set.add(task_family)
            selected_task_families.append(task_family)
        if slice_plan["episode_limit"] is not None and len(selected_entries) >= slice_plan["episode_limit"]:
            break

    if not selected_entries:
        raise RuntimeError(
            "MiniWoB slice selection removed every task. Relax --task-include / --task-exclude or choose a broader --slice-preset."
        )

    benchmark.env_args_list = [entry["env_args"] for entry in selected_entries]
    is_full_benchmark = len(selected_entries) == len(entries) and slice_plan["preset"] == "full"
    benchmark_label = "MiniWoB" if is_full_benchmark else f"MiniWoB Slice ({slice_plan['label']})"
    return benchmark, {
        "benchmark_label": benchmark_label,
        "is_full_benchmark": is_full_benchmark,
        "slice_preset": slice_plan["preset"],
        "slice_label": slice_plan["label"],
        "slice_description": slice_plan["description"],
        "task_include": include_patterns,
        "task_exclude": exclude_patterns,
        "seeds_per_task": slice_plan["seeds_per_task"],
        "task_limit": slice_plan["task_limit"],
        "episode_limit": slice_plan["episode_limit"],
        "original_env_count": len(entries),
        "selected_env_count": len(selected_entries),
        "selected_task_count": len(selected_task_families),
        "selected_task_families": selected_task_families,
        "selected_task_names": [entry["task_name"] for entry in selected_entries],
    }


def resolve_slice_plan(args: argparse.Namespace) -> Dict[str, Any]:
    preset = SLICE_PRESETS[args.slice_preset]
    task_include = list(preset["task_include"]) + list(args.task_include or [])
    task_exclude = list(preset["task_exclude"]) + list(args.task_exclude or [])
    seeds_per_task = preset["seeds_per_task"] if args.seeds_per_task is None else args.seeds_per_task
    task_limit = preset["task_limit"] if args.task_limit is None else args.task_limit
    episode_limit = preset["episode_limit"] if args.episode_limit is None else args.episode_limit
    is_customized = bool(args.task_include or args.task_exclude or args.seeds_per_task is not None or args.task_limit is not None or args.episode_limit is not None)
    label = args.slice_preset if not is_customized else f"{args.slice_preset}+custom"
    return {
        "preset": args.slice_preset,
        "label": label,
        "description": str(preset["description"]),
        "task_include": task_include,
        "task_exclude": task_exclude,
        "seeds_per_task": seeds_per_task,
        "task_limit": task_limit,
        "episode_limit": episode_limit,
    }


def build_output_notes(benchmark_meta: Mapping[str, Any]) -> List[str]:
    notes = [
        "This run uses AgentLab's MiniWoB benchmark flow and produces a real study_id.",
    ]
    if benchmark_meta.get("is_full_benchmark", True):
        notes.extend(
            [
                "Review README metadata fields and benchmark_specific / benchmark_tuned / reproducible before submitting.",
                "The generated miniwob.json is a draft submission row for the BrowserGym leaderboard.",
            ]
        )
    else:
        notes.extend(
            [
                f"This run only covers a filtered MiniWoB slice (`{benchmark_meta.get('slice_label', 'slice')}`).",
                "Use it as a cheaper real AgentLab check, not as a final full MiniWoB leaderboard submission.",
            ]
        )
    return notes


def print_selected_slice(benchmark_meta: Mapping[str, Any]) -> None:
    payload = {
        "benchmark": benchmark_meta.get("benchmark_label"),
        "slice_preset": benchmark_meta.get("slice_preset"),
        "slice_label": benchmark_meta.get("slice_label"),
        "selected_task_count": benchmark_meta.get("selected_task_count"),
        "selected_env_count": benchmark_meta.get("selected_env_count"),
        "selected_task_families": benchmark_meta.get("selected_task_families"),
        "selected_task_names": benchmark_meta.get("selected_task_names"),
    }
    print(json.dumps(payload, indent=2))


if __name__ == "__main__":
    main()