#!/usr/bin/env python3
"""Extend the sequential chains of a completed run so budgets can be compared.

The main experiment stops at K=4, which is enough to see whether revision helps
but not enough to answer the question the field actually cares about:

    given a fixed inference budget, is it better spent going wide or going deep?

Comparing N=16 against K=4 does not answer that, because they are not the same
budget. This continues each recorded chain to K=16 so parallel and sequential can
be read at matched budgets B in {1,2,4,8,16}.

The chains are *continued*, not regenerated. Each stored round has its response
text, so the exact conversation the original run built is reconstructed and
extended, which keeps rounds 1-4 identical to what the main experiment already
reported rather than re-rolling them.

Also records per-call latency, which the main run does not. Tokens measure what
you pay; latency measures what you wait, and the two arms differ on the second in
a way tokens alone hide: 16 samples can be issued at once, while 16 revisions
cannot start until the previous one finishes.

    python3 tools/extend_sequential.py <cache_key> [--rounds=16]
"""

from __future__ import annotations

import asyncio
import json
import sys
import time
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(REPO_ROOT / "backend"))

from app.env_fix import repair_ssl_env  # noqa: E402

repair_ssl_env()

import httpx  # noqa: E402
import ttc_core as core  # noqa: E402

CACHE_DIR = REPO_ROOT / "run_cache"


def rebuild_messages(question: str, chain: list[dict]) -> list[dict]:
    """Reconstruct the conversation the original run held for this question.

    Mirrors ttc_core.run_sequential_rounds exactly: system, question, the
    baseline answer, then alternating REVISE_PROMPT and response for each
    recorded round.
    """
    messages = [
        {"role": "system", "content": core.SYSTEM_PROMPT},
        {"role": "user", "content": question},
        {"role": "assistant", "content": chain[0]["text"]},
    ]
    for entry in chain[1:]:
        messages.append({"role": "user", "content": core.REVISE_PROMPT})
        messages.append({"role": "assistant", "content": entry["text"]})
    return messages


async def extend_one(client, sem, question, chain, gold, target, api_key, model):
    messages = rebuild_messages(question, chain)
    cumulative = chain[-1]["tokens"]
    out = []
    for r in range(len(chain) + 1, target + 1):
        messages.append({"role": "user", "content": core.REVISE_PROMPT})
        started = time.perf_counter()
        text, tokens = await core.call_model(client, messages, 0.0, sem, api_key, model)
        latency = time.perf_counter() - started
        cumulative += tokens
        messages.append({"role": "assistant", "content": text})
        answer = core.extract_answer(text)
        out.append({
            "round": r,
            "answer": answer,
            "tokens": cumulative,
            "call_tokens": tokens,
            "correct": core.is_correct(answer, gold),
            "latency_s": round(latency, 3),
        })
    return out


async def main() -> int:
    args = [a for a in sys.argv[1:] if not a.startswith("--")]
    target = next((int(a.split("=")[1]) for a in sys.argv[1:] if a.startswith("--rounds=")), 16)
    if not args:
        print(__doc__)
        return 2
    path = CACHE_DIR / f"{args[0]}.json"
    if not path.exists():
        print(f"No cached run at {path}", file=sys.stderr)
        return 2

    run = json.loads(path.read_text())
    model = run["config"]["model"]
    existing = run["config"]["sequential_rounds"]
    api_key = [l.split("=", 1)[1].strip() for l in open(REPO_ROOT / "backend" / ".env")
               if l.startswith("TOGETHER_API_KEY=")][0]

    n = len(run["questions"])
    new_rounds = target - existing
    print(f"model      : {model}")
    print(f"questions  : {n}")
    print(f"rounds     : K={existing} already recorded, extending to K={target}")
    print(f"calls      : {n * new_rounds}\n")
    if new_rounds <= 0:
        print("Nothing to extend.")
        return 1

    sem = asyncio.Semaphore(core.MAX_CONCURRENCY)
    done = {"n": 0}
    async with httpx.AsyncClient(timeout=120) as client:
        async def one(i):
            extra = await extend_one(
                client, sem, run["questions"][i]["question"], run["sequential"][i],
                run["questions"][i]["gold"], target, api_key, model)
            done["n"] += 1
            if done["n"] % 10 == 0:
                print(f"  {done['n']}/{n} chains extended", flush=True)
            return extra
        extended = await asyncio.gather(*[one(i) for i in range(n)])

    out = REPO_ROOT / "fixtures" / f"sequential_ext_{args[0]}.json"
    out.write_text(json.dumps({
        "cache_key": args[0], "run_id": run["run_id"], "model": model,
        "from_round": existing + 1, "to_round": target,
        "chains": [
            [{"round": e["round"], "answer": e["answer"], "tokens": e["tokens"],
              "call_tokens": e["call_tokens"], "correct": e["correct"],
              "latency_s": e["latency_s"]} for e in chain]
            for chain in extended
        ],
    }, indent=2))
    print(f"\nwritten -> {out.relative_to(REPO_ROOT)}")

    hard = [i for i, b in enumerate(run["baselines"]) if not b["correct"]]
    for k in (4, 8, 12, 16):
        acc = 0
        for i in hard:
            full = run["sequential"][i] + extended[i]
            entry = next((e for e in full if e["round"] == k), None)
            acc += bool(entry and entry["correct"])
        print(f"  hard subset accuracy at K={k:2}: {acc}/{len(hard)}")
    return 0


if __name__ == "__main__":
    raise SystemExit(asyncio.run(main()))