"""Time one real browser step (recorded in finetune/out/dagger_cases.jsonl) end to end inside the server process.
    python apps/profile_step.py <checkpoint dir>"""
import json, os, sys
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))); sys.path.insert(0, "/home/ckl/projects/S/laya/finetune")
os.environ.setdefault("LAYA_FMT", "v2")
from common import get_agent
from common_ft import build_request
from fast_batch import profile_step
agent = get_agent(sys.argv[1])
if agent.cfg.get("head_max_len_train"): agent.cfg["head_max_len"] = agent.cfg["head_max_len_train"]
cases = [json.loads(l) for l in open("/home/ckl/projects/S/laya/finetune/out/dagger_cases.jsonl")]
for c in cases[:4]:
    page = c["page_obj"]; state, questions, _, _ = build_request(page, c["goal"], c.get("history", []))
    n_opts = sum(len(q["criteria"]) for q in questions.values())
    r = profile_step(agent, state, questions)
    print(f"{len(questions)} questions / {n_opts:3d} options / {r['tokens']:4d} tokens: agent.predict {r['predict_ms']:6.1f} ms | fast {r['predict_fast_ms']:6.1f} ms "
          f"(tokenize {r['tokenize_ms']:.1f} + forward {r['forward_ms']:.1f} + post {r['post_ms']:.1f})")