#!/usr/bin/env python3
"""
prepare_data.py — one-time data extractor for the Atlassian Cloud RL "trajectory
playground" (a representational HuggingFace Static Space).

It reads REAL recorded artifacts from the sibling `atlassian_env/` project and emits a
single self-contained `data/playground.json` that the static page fetches at runtime.
`trajectories/` is gitignored in the env, so committing this JSON is what makes the Space
self-contained.

Run it with the repo's uv project Python (recommended — it also lets the extractor drive
the real environment to regenerate proper get -> act -> get gold reference trajectories):

    cd ../atlassian_env && uv run python ../atlassian_space/prepare_data.py

It also runs without the env deps (pure standard library), in which case the recorded
gold trajectories are used as-is:

    python3 prepare_data.py

Sources (all under ../atlassian_env, read-only):
  - tasks/registry.json + tasks/<id>.json        -> task metadata + the user instruction
  - tasks/gold/<id>.json                          -> gold reference call sequence
  - rubrics/<id>.py (module docstring)            -> human-readable rubric description
  - rollout/prompts.py                            -> the exact agent system prompt
  - tasks/gold/<id>.json (calls)                  -> driven through the real env to make the
                                                     gold reference trajectories (get->act->get)
  - trajectories/reward-spread-sweep/             -> the 18x7 model sweep over the hardened
    (+ sweep_summary.json)                           tasks/rubrics, with recorded reasoning

Nothing here mutates the environment source. The disposable work DB is the env's own
scratch file. Output is deterministic given the same inputs.
"""
from __future__ import annotations

import ast
import json
import os
import sys
import tempfile
from collections import Counter, defaultdict
from pathlib import Path

# --------------------------------------------------------------------------------------
# Paths
# --------------------------------------------------------------------------------------
HERE = Path(__file__).resolve().parent                       # .../atlassian_space
REPO_ROOT = HERE.parent                                       # .../atlassian_env_rleaas
ENV_ROOT = Path(os.environ.get("ENV_ROOT", REPO_ROOT / "atlassian_env")).resolve()

TASKS_DIR = ENV_ROOT / "tasks"
GOLD_DIR = TASKS_DIR / "gold"
RUBRICS_DIR = ENV_ROOT / "rubrics"
TRAJ_ROOT = ENV_ROOT / "trajectories"
TRAJ_SWEEP = TRAJ_ROOT / "reward-spread-sweep"
SWEEP_SUMMARY = TRAJ_SWEEP / "sweep_summary.json"

OUT_PATH = HERE / "data" / "playground.json"

if str(ENV_ROOT) not in sys.path:
    sys.path.insert(0, str(ENV_ROOT))

REASONING_CAP = 1600   # cap recorded reasoning/message length (one run loops to ~80k chars)

# --------------------------------------------------------------------------------------
# Constants sourced from the env
# --------------------------------------------------------------------------------------
TOOL_COUNT = 102
SHAPING_SIGNALS = [
    {"name": "R_PROGRESS", "value": 0.1, "detail": "a successful mutating call that changed semantic state"},
    {"name": "R_NEUTRAL", "value": 0.0, "detail": "a successful read, or a mutating call that changed nothing"},
    {"name": "R_ERROR", "value": -0.2, "detail": "tool returned a 4xx/5xx (malformed / contract-violating)"},
    {"name": "R_REDUNDANT", "value": -0.25, "detail": "identical call, no intervening state change (farming)"},
    {"name": "R_UNAVAILABLE", "value": -0.5, "detail": "called a tool not in the task's available_tools"},
]
REWARD_MODEL = {"partial_cap": 0.9, "gate_cap": 0.3, "partial_credit_cap_config": 0.99,
                "shaping_signals": SHAPING_SIGNALS}

SYSTEM_PROMPT_FALLBACK = (
    "You are an assistant operating an Atlassian Cloud workspace (Jira and Confluence) "
    "through a set of tools. Accomplish the user's task using ONLY the tools provided to "
    "you; do not assume any capability a tool does not expose.\n"
    "\n"
    "Guidelines:\n"
    "- Inspect state with read tools before mutating, and verify the result afterward when "
    "it matters.\n"
    "- Call one or more tools per turn. Pass arguments exactly as each tool's schema "
    "requires.\n"
    "- If a tool returns an error (a 4xx/5xx status), read the message and adjust your next "
    "call instead of repeating the same one.\n"
    "- When the task is fully accomplished (or you have determined it cannot be done), call "
    "`finish` with a short `final_response` summarizing the outcome. If the task asks "
    "you to report a specific value, include that value verbatim in `final_response`.\n"
)

MODEL_META = {
    "gold":              {"label": "Gold (reference)",  "vendor": "reference",  "color": "#6366f1"},
    "claude-sonnet-4.6": {"label": "claude-sonnet-4.6", "vendor": "anthropic",  "color": "#d97757"},
    "gpt-5.5":           {"label": "gpt-5.5",           "vendor": "openai",     "color": "#10a37f"},
    "gemini-3.5-flash":  {"label": "gemini-3.5-flash",  "vendor": "google",     "color": "#4285f4"},
    "deepseek-v4-pro":   {"label": "deepseek-v4-pro",   "vendor": "deepseek",   "color": "#6d4dfc"},
    "kimi-k2.6":         {"label": "kimi-k2.6",         "vendor": "moonshotai", "color": "#e0457b"},
    "qwen3.7-plus":      {"label": "qwen3.7-plus",      "vendor": "qwen",       "color": "#a855f7"},
    "glm-5.2":           {"label": "glm-5.2",           "vendor": "z-ai",       "color": "#f59e0b"},
}
RUN_ORDER = {"gold": 0, "claude-sonnet-4.6": 1, "gpt-5.5": 2, "gemini-3.5-flash": 3,
             "deepseek-v4-pro": 4, "kimi-k2.6": 5, "qwen3.7-plus": 6, "glm-5.2": 7}
LEADERBOARD_MODELS = ["claude-sonnet-4.6", "gpt-5.5", "gemini-3.5-flash",
                      "deepseek-v4-pro", "kimi-k2.6", "qwen3.7-plus", "glm-5.2"]


# --------------------------------------------------------------------------------------
# Helpers
# --------------------------------------------------------------------------------------
def load_json(path: Path):
    with path.open(encoding="utf-8") as fh:
        return json.load(fh)


def read_jsonl(path: Path):
    out = []
    with path.open(encoding="utf-8") as fh:
        for line in fh:
            line = line.strip()
            if line:
                out.append(json.loads(line))
    return out


def cap(s):
    if not s:
        return s
    return s if len(s) <= REASONING_CAP else s[:REASONING_CAP].rstrip() + " … (truncated)"


def scenario_name(task_id: str) -> str:
    stem = task_id.split("-", 1)[1] if "-" in task_id else task_id
    return stem.replace("-", " ").title()


def rubric_description(task_id: str) -> str:
    path = RUBRICS_DIR / f"{task_id}.py"
    if not path.exists():
        return ""
    try:
        return (ast.get_docstring(ast.parse(path.read_text(encoding="utf-8"))) or "").strip()
    except SyntaxError:
        return ""


def rubric_scheme(task_id: str) -> str:
    path = RUBRICS_DIR / f"{task_id}.py"
    if path.exists() and 'scheme="multiplicative"' in path.read_text(encoding="utf-8"):
        return "multiplicative"
    return "weighted_sum"


def resolve_system_prompt() -> str:
    try:
        from rollout.prompts import SYSTEM_PROMPT  # type: ignore
        return SYSTEM_PROMPT
    except Exception:
        return SYSTEM_PROMPT_FALLBACK


def model_id_from_header(model_raw: str, kind: str) -> str:
    if model_raw == "scripted-gold" or kind == "gold":
        return "gold"
    return model_raw.rsplit("/", 1)[-1]


def parse_trajectory(path: Path, kind: str, end_reason: str | None = None) -> dict:
    records = read_jsonl(path)
    header = next((r for r in records if r.get("record") == "header"), {})
    terminal = next((r for r in records if r.get("record") == "terminal"), {})
    step_records = [r for r in records if r.get("record") == "step"]
    assistant_records = [r for r in records if r.get("record") == "assistant"]

    # Map each tool_call_id -> the assistant turn that issued it (for reasoning/content).
    turns = {}
    tcid_to_turn = {}
    for a in assistant_records:
        ti = a.get("turn_index")
        turns[ti] = {"reasoning": cap(a.get("reasoning")), "content": cap(a.get("content"))}
        for tcid in (a.get("tool_call_ids") or []):
            tcid_to_turn[tcid] = ti

    steps = []
    for s in step_records:
        out = s.get("output")
        ti = tcid_to_turn.get(s.get("tool_call_id"))
        turn = turns.get(ti, {})
        steps.append({
            "step_index": s.get("step_index"),
            "tool_name": s.get("tool_name"),
            "tool_call_id": s.get("tool_call_id"),
            "args": s.get("args"),
            "output": out,
            "status": out.get("status") if isinstance(out, dict) else None,
            "reward": s.get("reward"),
            "turn_index": ti,
            "reasoning": turn.get("reasoning"),
            "content": turn.get("content"),
        })

    # The model's closing message: the last assistant turn that issued no tool calls.
    final_msg = None
    for a in assistant_records:
        if not (a.get("tool_call_ids") or []) and a.get("content"):
            final_msg = a.get("content")
    final_response = terminal.get("final_response") or cap(final_msg)

    model_raw = header.get("model", "")
    model_id = model_id_from_header(model_raw, kind)
    meta = MODEL_META.get(model_id, {"label": model_id, "vendor": "model", "color": "#64748b"})
    return {
        "run_id": path.stem.rsplit("__", 1)[-1],
        "file": path.name,
        "kind": kind,
        "model_raw": model_raw,
        "model_id": model_id,
        "model_label": meta["label"],
        "vendor": meta["vendor"],
        "color": meta["color"],
        "order": RUN_ORDER.get(model_id, 9),
        "seed": header.get("seed"),
        "start_ts": header.get("start_ts"),
        "config_hash": header.get("config_hash"),
        "final_reward": terminal.get("final_reward"),
        "passed": terminal.get("passed"),
        "n_steps": terminal.get("n_steps", len(steps)),
        "terminated_reason": end_reason or terminal.get("terminated_reason"),
        "final_response": final_response,
        "components": terminal.get("components", []),
        "steps": steps,
        "note": "",
    }


def rebuild_gold(task_id: str, gold_calls, final_response=None) -> dict | None:
    """Drive the task's gold solution calls through the REAL environment to produce a
    genuine gold reference trajectory (the gold solutions are already proper
    get -> act -> get sequences). Returns None only if the env can't be driven."""
    try:
        from orchestration.env import RLEnv
    except Exception:
        return None
    tmp = Path(tempfile.mkdtemp(prefix="goldgen_"))
    try:
        env = RLEnv(task_id, env_root=ENV_ROOT, model="scripted-gold", trajectories_dir=tmp)
        env.reset(seed=42)
        for c in gold_calls:
            env.step({"name": c["name"], "args": c.get("args", {})})
        env.finish(final_response=final_response)
        path = env.trajectory_path
        env.close()
        return parse_trajectory(Path(path), "gold")
    except Exception as exc:
        print(f"  ! gold rebuild failed for {task_id}: {exc!r}")
        return None


# --------------------------------------------------------------------------------------
# Build
# --------------------------------------------------------------------------------------
def build():
    if not ENV_ROOT.exists():
        sys.exit(f"ENV_ROOT not found: {ENV_ROOT}\nSet ENV_ROOT=/path/to/atlassian_env and re-run.")

    registry = load_json(TASKS_DIR / "registry.json")["tasks"]
    runs_by_task: dict[str, list[dict]] = defaultdict(list)

    # gold reference trajectories: drive each task's gold solution through the real env
    gold_ok, gold_bad = 0, []
    for entry in registry:
        tid = entry["task_id"]
        gp = GOLD_DIR / f"{tid}.json"
        if not gp.exists():
            continue
        g = load_json(gp)
        run = rebuild_gold(tid, g.get("calls", []), g.get("final_response"))
        if run is None:  # env not drivable -> fall back to any recorded scripted gold
            rec = next((p for p in TRAJ_ROOT.glob(f"*__{tid}__*.jsonl")), None)
            run = parse_trajectory(rec, "gold") if rec else None
        if run is None:
            continue
        runs_by_task[tid].append(run)
        if run["final_reward"] == 1.0 and run["passed"]:
            gold_ok += 1
        else:
            gold_bad.append((tid, run["final_reward"]))

    # model sweep: one run per (task, model). sweep_summary.json is the authority for the
    # canonical reward + end reason; pick the trajectory file that matches it (a few cells
    # have a second attempt), skipping any incomplete file with no terminal record.
    summary = {}
    if SWEEP_SUMMARY.exists():
        for e in load_json(SWEEP_SUMMARY):
            summary[(e["task"], e["model"])] = e

    groups: dict[tuple, list[dict]] = defaultdict(list)
    if TRAJ_SWEEP.exists():
        for path in sorted(TRAJ_SWEEP.glob("*.jsonl")):
            tid = path.stem.split("__")[1] if "__" in path.stem else None
            if tid is None:
                continue
            run = parse_trajectory(path, "sweep")
            if run["final_reward"] is None:      # incomplete attempt (no terminal record)
                continue
            groups[(tid, run["model_id"])].append(run)

    sweep_n = 0
    for (tid, mid), cand in groups.items():
        s = summary.get((tid, mid))
        if s is not None:                        # prefer the file matching the canonical reward
            matches = [r for r in cand if abs((r["final_reward"] or 0) - s["reward"]) < 1e-9]
            cand = matches or cand
        run = max(cand, key=lambda r: r["start_ts"] or "")
        if s is not None:
            run["terminated_reason"] = s.get("end")
        runs_by_task[tid].append(run)
        sweep_n += 1

    # ---- per-task records ----------------------------------------------------------
    tasks = []
    for entry in registry:
        task_id = entry["task_id"]
        task = load_json(TASKS_DIR / entry["file"])

        gold = None
        gp = GOLD_DIR / f"{task_id}.json"
        if gp.exists():
            g = load_json(gp)
            gold = {"intended_state": g.get("intended_state", ""), "calls": g.get("calls", []),
                    "n_calls": len(g.get("calls", []))}

        runs = sorted(runs_by_task.get(task_id, []), key=lambda r: (r["order"], r["start_ts"] or ""))
        gold_run = next((r for r in runs if r["kind"] == "gold"), None)
        comp_source = (gold_run["components"] if gold_run and gold_run["components"]
                       else (runs[0]["components"] if runs else []))
        rubric_components = [{"name": c.get("name"), "weight": c.get("weight"),
                              "gate": bool(c.get("gate")), "detail": c.get("detail", "")} for c in comp_source]

        model_runs = [r for r in runs if r["model_id"] != "gold"]
        best_reward = max((r["final_reward"] for r in model_runs if r["final_reward"] is not None), default=None)
        models_run = sorted({r["model_id"] for r in model_runs}, key=lambda m: RUN_ORDER.get(m, 9))

        tasks.append({
            "task_id": task_id,
            "name": scenario_name(task_id),
            "title": task.get("title", entry.get("title", task_id)),
            "product": task.get("product", entry.get("product")),
            "difficulty": task.get("difficulty"),
            "instruction": task.get("instruction", ""),
            "available_tools": task.get("available_tools", []),
            "initial_state_notes": task.get("initial_state_notes", ""),
            "max_steps": task.get("max_steps"),
            "tags": task.get("tags", []),
            "rubric_path": task.get("rubric", f"rubrics/{task_id}.py"),
            "rubric": {"description": rubric_description(task_id), "aggregation": rubric_scheme(task_id),
                       "components": rubric_components},
            "gold": gold,
            "runs": runs,
            "has_live": bool(model_runs),
            "best_reward": best_reward,
            "models_run": models_run,
            "gold_steps": gold_run["n_steps"] if gold_run else (gold["n_calls"] if gold else None),
        })

    # ---- leaderboard over all swept tasks -----------------------------------------
    swept_task_ids = [t["task_id"] for t in tasks if t["has_live"]]

    def reward_for(tid, model_id):
        for r in runs_by_task.get(tid, []):
            if r["model_id"] == model_id and r["final_reward"] is not None:
                return r["final_reward"]
        return None

    leaderboard_rows = []
    for mid in LEADERBOARD_MODELS:
        per_task = {t: reward_for(t, mid) for t in swept_task_ids}
        present = [v for v in per_task.values() if v is not None]
        leaderboard_rows.append({
            "model_id": mid, "label": MODEL_META[mid]["label"], "color": MODEL_META[mid]["color"],
            "kind": "sweep", "order": RUN_ORDER.get(mid, 9), "per_task": per_task,
            "mean_reward": round(sum(present) / len(present), 4) if present else None,
            "pass_count": sum(1 for v in present if v >= 1.0), "task_count": len(present),
        })

    # ---- meta ----------------------------------------------------------------------
    diff_counts = Counter(t["difficulty"] for t in tasks)
    prod_counts = Counter(t["product"] for t in tasks)
    meta = {
        "env_id": "atlassian-cloud-rl",
        "title": "Atlassian Cloud RL Environment",
        "subtitle": "Tool-use RL environment cloning Atlassian Cloud (Jira v3 + Agile + Confluence v2/v1) over a SQLite store.",
        "counts": {"tasks": len(tasks), "tools": TOOL_COUNT,
                   # the swept model trajectories (gold reference replays are counted separately)
                   "trajectories": sum(1 for t in tasks for r in t["runs"] if r["model_id"] != "gold"),
                   "products": ["jira", "confluence"]},
        "difficulty_counts": dict(diff_counts),
        "product_counts": dict(prod_counts),
        "system_prompt": resolve_system_prompt(),
        "reward_model": REWARD_MODEL,
        "models": [{"id": mid, **MODEL_META[mid]} for mid in ["gold"] + LEADERBOARD_MODELS],
        "swept_task_ids": swept_task_ids,
        "source_note": "All values are read verbatim from the recorded JSONL trajectories and task/rubric files in atlassian_env/; gold reference runs are produced by driving the real environment.",
    }

    payload = {"meta": meta, "leaderboard": {"swept_task_ids": swept_task_ids, "rows": leaderboard_rows}, "tasks": tasks}
    OUT_PATH.parent.mkdir(parents=True, exist_ok=True)
    with OUT_PATH.open("w", encoding="utf-8") as fh:
        json.dump(payload, fh, ensure_ascii=False, separators=(",", ":"))

    # ---- console summary -----------------------------------------------------------
    print(f"Wrote {OUT_PATH.relative_to(REPO_ROOT)}  ({OUT_PATH.stat().st_size/1024:.0f} KB)")
    print(f"  tasks: {len(tasks)} ({dict(prod_counts)}, {dict(diff_counts)})  | gold passing: {gold_ok}/{len(tasks)}")
    if gold_bad:
        print(f"  ! gold not at 1.0: {gold_bad}")
    print(f"  sweep runs: {sweep_n}  | total runs: {meta['counts']['trajectories']}  | swept tasks: {len(swept_task_ids)}  | models: {len(LEADERBOARD_MODELS)}")
    gmin = min((t["gold_steps"] for t in tasks if t["gold_steps"] is not None), default=None)
    print(f"  min gold steps across tasks: {gmin}")
    for r in leaderboard_rows:
        mr = "n/a" if r["mean_reward"] is None else f"{r['mean_reward']:.3f}"
        print(f"    {r['label']:20s} mean={mr}  pass={r['pass_count']}/{r['task_count']}")


if __name__ == "__main__":
    build()