"""Same-day Jev baseline on LocalLLaMA/typed-decisions test (evaluation only).

One request per case with all of its questions, as the dataset's leaderboard scored Jev.
Jev's answers are written to their own file and are never used for training, calibration or
model selection (TypeSafe MCA 2.3(b)).

Needs TYPESAFE_API_KEY in the environment.
"""

import argparse
import concurrent.futures as cf
import json
import os
import random
import time

import requests
from datasets import load_dataset

from td_metrics import parsed, score_cases

URL = "https://api.typesafe.ai/v1/systemone"


def call(session, key, model, row):
    body = {"model": model, "state": parsed(row["state"]), "questions": parsed(row["questions"])}
    for attempt in range(10):
        try:
            response = session.post(
                URL, json = body, headers = {"Authorization": f"Bearer {key}"}, timeout = 180
            )
        except requests.RequestException as exc:
            error = repr(exc)
        else:
            if response.status_code == 200:
                return response.json()
            error = f"HTTP {response.status_code}: {response.text[:300]}"
            if response.status_code not in (408, 409, 425, 429, 500, 502, 503, 504, 529):
                raise RuntimeError(error)
        time.sleep(min(60.0, 2 ** attempt) + random.random())
    raise RuntimeError(f"{row['id']}: gave up: {error}")


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default = "jev-latest")
    parser.add_argument("--out", default = "/content/out/jev_td_test.jsonl")
    parser.add_argument("--workers", type = int, default = 8)
    parser.add_argument("--rows", default = None, help = "jsonl of id/workflow/state/questions/gold rows")
    args = parser.parse_args()
    key = os.environ["TYPESAFE_API_KEY"]
    if args.rows:
        rows = [json.loads(line) for line in open(args.rows, encoding = "utf-8") if line.strip()]
    else:
        rows = list(load_dataset("LocalLLaMA/typed-decisions", "all", split = "test"))
    os.makedirs(os.path.dirname(args.out), exist_ok = True)
    done = {}
    if os.path.exists(args.out):
        for line in open(args.out, encoding = "utf-8"):
            if line.strip():
                record = json.loads(line)
                done[record["id"]] = record
    todo = [row for row in rows if row["id"] not in done]
    print(f"jev: {len(done)} cached, {len(todo)} to request", flush = True)
    start = time.time()
    session = requests.Session()
    with open(args.out, "a", encoding = "utf-8") as file, cf.ThreadPoolExecutor(args.workers) as pool:
        futures = {pool.submit(call, session, key, args.model, row): row for row in todo}
        for i, future in enumerate(cf.as_completed(futures), 1):
            row = futures[future]
            response = future.result()
            record = {
                "id": row["id"],
                "workflow": row["workflow"],
                "model": response.get("model"),
                "usage": response.get("usage"),
                "pred": response["answers"],
            }
            done[row["id"]] = record
            file.write(json.dumps(record) + "\n")
            file.flush()
            if i % 50 == 0:
                print(f"jev: {i}/{len(todo)} in {time.time() - start:.0f}s", flush = True)
    cases = [
        {
            "id": row["id"],
            "workflow": row["workflow"],
            "questions": row["questions"],
            "gold": row["gold"],
            "pred": done[row["id"]]["pred"],
        }
        for row in rows
    ]
    metrics = score_cases(cases)
    metrics["model_reported"] = sorted({done[r["id"]]["model"] for r in rows})
    metrics["input_tokens"] = sum((done[r["id"]].get("usage") or {}).get("input_tokens", 0) for r in rows)
    with open(args.out.replace(".jsonl", "_metrics.json"), "w", encoding = "utf-8") as file:
        json.dump(metrics, file, indent = 1)
    print("JEV_METRICS " + json.dumps(metrics), flush = True)


if __name__ == "__main__":
    main()