"""Stage 2: turn graded runs into routing labels (chat-format SFT data) with a held-out test split.

Label policy (cost-aware oracle):
  * cheapest routable agent that solved the task, counting a win only if the agent is reliable on that source
    (solve rate >= 70%). Wins by unreliable agents are mostly lucky single samples a router cannot predict: v1 sent
    gemma-3-1b 68 held-out tasks and only 53% were solved, and kimi-vl solves just 63% of text tasks that >= 4 other
    agents solve (lightning 95%, glm 98%). gemma-3-1b is excluded as a target outright;
  * if no routable agent solved it -> ensemble of the capable agents among glm / qwen / ultra (aggregated by the
    last one) when at least two can take the task, otherwise the strongest capable agent alone.
The router sees `router_view(...)`: a size / image header plus the start and end of the task.
"""
import json, random
from collections import Counter
from datasets import Dataset, DatasetDict
from common import AGENTS, POOL, ROUTER_SYSTEM, STRONGEST, capable, est_tokens, router_view

random.seed(0)
ROUTABLE = [a for a in AGENTS if a != "gemma-3-1b"]
ENSEMBLE_POOL = ["glm-5.3-flash", "qwen3.8-27b", STRONGEST]
RELIABLE = 0.70
SOLVE_RATE = {}  # (source, agent) -> solve rate, filled in main()


def cost(res, name):
    return POOL[name]["cost"] * (res[name]["prompt_tokens"] + res[name]["completion_tokens"])


def can(task, agent):
    """Capability check; long prompts are stored truncated, so trust the recorded skip instead of re-estimating."""
    return not task["results"][agent].get("skipped") and capable(agent, "" if task.get("prompt_truncated")
                                                                 else task["prompt"], task["has_image"])


def label(task):
    res = task["results"]
    n_correct = sum(res[a]["correct"] for a in AGENTS)
    correct = [a for a in ROUTABLE if res[a]["correct"] and SOLVE_RATE[task["source"], a] >= RELIABLE]
    difficulty = "easy" if n_correct >= 4 else "medium" if n_correct >= 2 else "hard"
    if correct:
        return dict(strategy="single", agents=[min(correct, key=lambda a: cost(res, a))], difficulty=difficulty)
    ens = [a for a in ENSEMBLE_POOL if can(task, a)]
    if len(ens) >= 2 and not task["has_image"]:
        return dict(strategy="ensemble", agents=ens, difficulty="hard")
    fallback = [a for a in ROUTABLE if can(task, a)] or [STRONGEST]
    return dict(strategy="single", agents=[fallback[-1]], difficulty="hard")


def to_example(task):
    lab = label(task)
    target = json.dumps({"domain": task["domain"], "difficulty": lab["difficulty"],
                         "strategy": lab["strategy"], "agents": lab["agents"]})
    view = router_view(task["prompt"], task["has_image"], task["est_tokens"])
    return dict(id=task["id"], source=task["source"], domain=task["domain"], kind=task["kind"],
                gold=task["gold"], prompt=task["prompt"], has_image=task["has_image"],
                est_tokens=task["est_tokens"], prompt_truncated=bool(task.get("prompt_truncated")),
                router_input=view, route=target,
                agent_results=json.dumps({a: {k: task["results"][a].get(k) for k in
                                              ("correct", "completion_tokens", "latency", "skipped")}
                                          for a in AGENTS}),
                messages=[{"role": "system", "content": ROUTER_SYSTEM},
                          {"role": "user", "content": view},
                          {"role": "assistant", "content": target}])


def main():
    rows = [json.loads(l) for l in open("data/runs_v3.jsonl")]
    n_all = len(rows)
    rows = [r for r in rows if not any(v.get("error") for v in r["results"].values())]
    print(f"{n_all} runs, {n_all - len(rows)} dropped for transient API errors")
    for r in rows:
        r.setdefault("has_image", False)
        r.setdefault("est_tokens", est_tokens(r["prompt"]))
    for src in {r["source"] for r in rows}:
        sub = [r for r in rows if r["source"] == src]
        for a in AGENTS:
            SOLVE_RATE[src, a] = sum(r["results"][a]["correct"] for r in sub) / len(sub)
    random.shuffle(rows)
    ex = [to_example(r) for r in rows]
    n_test = min(max(200, len(ex) // 10), len(ex) // 5)
    dd = DatasetDict(train=Dataset.from_list(ex[n_test:]), test=Dataset.from_list(ex[:n_test]))
    dd.save_to_disk("data/shaktimaan_routing")
    print(dd)
    routes = Counter(json.loads(e["route"])["strategy"] + ":" + ",".join(json.loads(e["route"])["agents"])
                     for e in ex)
    print("label distribution:", *routes.most_common(), sep="\n  ")
    for src in sorted({r["source"] for r in rows}):
        sub = [r for r in rows if r["source"] == src]
        print(f"{src:10s} n={len(sub):4d} " + " ".join(
            f"{a}={sum(r['results'][a]['correct'] for r in sub) / len(sub):.2f}" for a in AGENTS))


if __name__ == "__main__":
    main()