# Strict BIRD execution accuracy: set(pred rows) == set(gold rows). usage: python score.py preds.jsonl db_root
import json, os, sqlite3, sys, collections, time
from concurrent.futures import ProcessPoolExecutor, TimeoutError as TO

HERE = os.path.dirname(os.path.abspath(__file__))
preds_path, db_root = sys.argv[1], sys.argv[2]
gold = {d["question_id"]: d for d in (json.loads(l) for l in open(os.path.join(HERE, "mini_dev_sqlite.jsonl"), encoding="utf-8"))}
preds = {}
for l in open(preds_path, encoding="utf-8"):
    p = json.loads(l); preds[p["question_id"]] = p


def run(db, sql, limit=None):
    con = sqlite3.connect(f"file:{db}?mode=ro", uri=True)
    if limit:  # abort runaway queries inside the worker; a pool timeout alone leaves them running
        end = time.time() + limit
        con.set_progress_handler(lambda: time.time() > end, 100000)
    try:
        return set(con.execute(sql).fetchall())
    finally:
        con.close()


def judge(args):
    db, psql, gsql = args
    try:
        g = run(db, gsql)
    except Exception as e:
        return "gold_error"
    if not psql: return "no_sql"
    try:
        return "ok" if run(db, psql, 60) == g else "wrong"
    except sqlite3.OperationalError as e:
        return "timeout" if "interrupted" in str(e) else "exec_error"
    except Exception:
        return "exec_error"


if __name__ == "__main__":
    jobs = []
    for qid, g in gold.items():
        db = os.path.join(db_root, g["db_id"], g["db_id"] + ".sqlite")
        jobs.append((qid, (db, preds.get(qid, {}).get("sql", ""), g["SQL"])))
    res = {}
    with ProcessPoolExecutor(8) as ex:
        futs = {qid: ex.submit(judge, a) for qid, a in jobs}
        for qid, fu in futs.items():
            try: res[qid] = fu.result(timeout=600)  # the 60 s per-query limit is enforced inside the worker; this only guards a hung pool
            except TO: res[qid] = "timeout"
    c = collections.Counter(res.values())
    by = collections.defaultdict(lambda: [0, 0])
    for qid, v in res.items():
        d = gold[qid]["difficulty"]; by[d][1] += 1; by[d][0] += v == "ok"
    print(f"EX = {c['ok']}/{len(res)} = {100 * c['ok'] / len(res):.2f}%   ", dict(c))
    for d, (a, n) in sorted(by.items()): print(f"  {d:12s} {a}/{n} = {100 * a / n:.1f}%")
    json.dump(res, open(preds_path.replace(".jsonl", "_score.json"), "w"), indent=0)