"""Main run: N categories x 2 native Perplexity stacks. Saves raw jsonl."""
import json, sys, os, time, threading
import concurrent.futures as cf
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
import orlib
from cats300 import CATS

MODELS = ["perplexity/sonar", "perplexity/sonar-pro"]
OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "main_raw.jsonl")

done = set()
if os.path.exists(OUT):
    for line in open(OUT):
        try:
            r = json.loads(line)
            done.add((r["model"], r["cat"]))
        except Exception:
            pass
print("already done:", len(done))

jobs = [(m, c) for c in CATS for m in MODELS if (m, c) not in done]
print("jobs:", len(jobs))

lock = threading.Lock()
fh = open(OUT, "a")
t0 = time.time()
n = [0]

def run(j):
    m, c = j
    rec = {"model": m, "cat": c, "ts": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())}
    for att in range(3):
        try:
            r = orlib.ask(m, c)
            rec.update({"picks": orlib.parse(r["text"]), "citations": r["citations"],
                        "usage": r["usage"], "text": r["text"][:2000]})
            break
        except Exception as e:
            if att == 2:
                rec.update({"error": repr(e)[:200], "picks": [], "citations": [], "usage": {}})
            else:
                time.sleep(4 + 4 * att)
    with lock:
        fh.write(json.dumps(rec) + "\n"); fh.flush()
        n[0] += 1
        if n[0] % 25 == 0:
            print(n[0], len(jobs), round(time.time() - t0), flush=True)
    return rec

with cf.ThreadPoolExecutor(12) as ex:
    res = list(ex.map(run, jobs))
fh.close()
print("done", len(res), "parsed", sum(1 for r in res if r.get("picks")),
      "errors", sum(1 for r in res if r.get("error")), "secs", round(time.time() - t0))