"""Generate teacher trajectories for training samples: prompts -> teacher answers + states files,
kept in /root/states. The trainer reads these.

    .venv/bin/python gen_teacher.py --dataset drawn.jsonl --samples 32 --max-prompt 1500
"""
import argparse
import json
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

from eval_service import dataset, run, servers
from eval_service.run import _render, generate, load_tokenizer

ap = argparse.ArgumentParser()
ap.add_argument("--dataset", type=Path, default=Path("/root/drawn.jsonl"))
ap.add_argument("--teacher", type=Path, default=Path("/root/models/teacher"))
ap.add_argument("--samples", type=int, default=32)
ap.add_argument("--max-prompt", type=int, default=1500)
ap.add_argument("--workers", type=int, default=8)
ap.add_argument("--out", type=Path, default=Path("/root/train_samples.json"))
args = ap.parse_args()

samples = dataset.load(args.dataset).samples
tokenizer = load_tokenizer(args.teacher)
picked, prompts = [], []
for s in samples:
    p = _render(tokenizer, s.messages, s.tools)
    if len(p) <= args.max_prompt:
        picked.append(s)
        prompts.append(p)
        if len(picked) >= args.samples:
            break
print(f"{len(picked)} samples, prompts {min(map(len, prompts))}-{max(map(len, prompts))} tokens")

out = Path("/root/states")
out.mkdir(exist_ok=True)
records = []


def one(i, s, p):
    answer, tpath = generate(groups["teacher"], p, i, s.id)
    for _ in range(120):  # the connector writes states asynchronously
        if Path(tpath).exists():
            break
        time.sleep(0.5)
    keep = out / f"t{i}.safetensors"
    keep.write_bytes(Path(tpath).read_bytes())
    Path(str(keep) + ".lock").touch()
    return answer, str(keep)


try:
    groups = servers.serve({"teacher": args.teacher})
    with ThreadPoolExecutor(args.workers) as pool:
        futures = {i: pool.submit(one, i, s, p) for i, (s, p) in enumerate(zip(picked, prompts))}
        for i, fut in futures.items():
            answer, keep = fut.result()
            records.append({"i": i, "id": picked[i].id, "prompt": prompts[i], "answer": answer, "states": keep})
            print(f"{i+1}/{len(picked)} {picked[i].id} p{len(prompts[i])} a{len(answer)}", flush=True)
finally:
    servers.stop_all()

args.out.write_text(json.dumps(records))
print("saved", args.out)