#!/usr/bin/env python3
"""Build the HoloBench viewer data.

HoloBench ships 5 relational databases whose rows are each *verbalized* into a
sentence (`verbalized_text`) alongside the structured row (`table_row`), plus a
set of SQL-derived queries per database (`query/<db>.csv`). Gold answers are not
stored — they are produced by executing the SQL. This script:

  * Corpus  -> one shard per table (incl. the `dummy` distractor tables):
      corpus/index.json           light index (db, table, #rows, #cols, distractor)
      corpus/records/<key>.json   sampled rows (verbalized_text + table_row), capped
  * Eval    -> eval.jsonl : one record per query with the SQL, metadata, the
      tables it references, and the gold answer computed by running the SQL over
      the FULL table data via DuckDB (display rows capped, total row count kept).

Run from the viewer repo root:
    python scripts/build_data.py [--src <dataset-dir>] [--sample-rows 300] [--answer-rows 100]

Default --src is $DATA_ROOT/holobench (falls back to /tmp/holo_raw).
"""
import argparse
import ast
import csv
import glob
import json
import os

import duckdb
import pandas as pd
from sql_metadata import Parser

csv.field_size_limit(10 ** 9)

DBS = ["wine_1", "college_2", "flight_4", "store_1", "soccer_1"]
FLAG_COLS = ["Aggregation", "Max/Min", "Join", "Comparison", "Ranking"]


def load_table_rows(path):
    """Return (all_rows, structured_rows) where each all_rows item is
    {verbalized_text, table_row(dict)}; structured_rows drops empty table_row."""
    all_rows, structured = [], []
    with open(path) as fh:
        for r in csv.DictReader(fh):
            tr_raw = r.get("table_row") or "{}"
            try:
                tr = json.loads(tr_raw)
            except Exception:
                tr = {}
            item = {"verbalized_text": r.get("verbalized_text", ""), "table_row": tr}
            all_rows.append(item)
            if tr:
                structured.append(tr)
    return all_rows, structured


def to_bool(v):
    return str(v).strip().lower() == "true"


def main():
    root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    default_src = os.path.join(
        os.environ.get("DATA_ROOT", "/mnt/ramdisk/blobstore/timchen0618/data"), "holobench"
    )
    if not os.path.isdir(default_src):
        default_src = "/tmp/holo_raw"

    ap = argparse.ArgumentParser()
    ap.add_argument("--src", default=default_src)
    ap.add_argument("--sample-rows", type=int, default=300, help="corpus rows shown per table")
    ap.add_argument("--answer-rows", type=int, default=100, help="gold-answer rows shown per query")
    args = ap.parse_args()

    records_dir = os.path.join(root, "corpus", "records")
    os.makedirs(records_dir, exist_ok=True)
    for fn in os.listdir(records_dir):
        if fn.endswith(".json"):
            os.remove(os.path.join(records_dir, fn))

    index = []
    eval_rows = []

    for db in DBS:
        base = os.path.join(args.src, db)
        # --- load every table once (full data) ---
        full = {}   # table -> DataFrame (structured, for SQL)
        for path in sorted(glob.glob(os.path.join(base, "*.csv"))):
            table = os.path.basename(path)[:-4]
            all_rows, structured = load_table_rows(path)
            distractor = (table == "dummy") or (len(structured) == 0)
            cols = list(structured[0].keys()) if structured else []
            if structured:
                full[table] = pd.DataFrame(structured)

            key = f"{db}__{table}"
            sample = all_rows[: args.sample_rows]
            with open(os.path.join(records_dir, key + ".json"), "w") as f:
                json.dump({
                    "db_id": db, "table": table, "columns": cols,
                    "total_rows": len(all_rows), "distractor": distractor,
                    "shown_rows": len(sample), "rows": sample,
                }, f, ensure_ascii=False)
            index.append({
                "db_id": db, "table": table, "key": key,
                "columns": cols, "total_rows": len(all_rows), "distractor": distractor,
            })

        # --- queries + gold answers ---
        qpath = os.path.join(args.src, "query", f"{db}.csv")
        for i, q in enumerate(csv.DictReader(open(qpath))):
            sql = q["query"]
            try:
                refs = [t for t in Parser(sql).tables]
            except Exception:
                refs = []
            gold = {}
            con = duckdb.connect(":memory:")
            try:
                for name, df in full.items():
                    con.register(name, df)
                res = con.query(sql).df()
                total = len(res)
                shown = res.head(args.answer_rows)
                gold = {
                    "columns": list(res.columns),
                    "rows": json.loads(shown.to_json(orient="values")),
                    "total_rows": total,
                    "truncated": total > args.answer_rows,
                }
            except Exception as e:
                gold = {"error": str(e)}
            finally:
                con.close()

            try:
                qtype = ast.literal_eval(q.get("query_type") or "[]")
            except Exception:
                qtype = [q.get("query_type")]
            eval_rows.append({
                "id": f"{db}-{i}",
                "db_id": db,
                "question": q["question"],
                "query": sql,
                "difficulty": q.get("difficulty"),
                "query_type": qtype,
                "flags": [c for c in FLAG_COLS if to_bool(q.get(c))],
                "sql_components": [q.get("SQL_COMPONENTS_1"), q.get("SQL_COMPONENTS_2"),
                                   q.get("Other Complexity Factors")],
                "tables": refs,
                "gold_answer": gold,
            })
        print(f"{db}: {len([r for r in eval_rows if r['db_id']==db])} queries, "
              f"{len([1 for x in index if x['db_id']==db])} tables")

    index.sort(key=lambda d: (d["db_id"], d["distractor"], d["table"]))
    with open(os.path.join(root, "corpus", "index.json"), "w") as f:
        json.dump(index, f, ensure_ascii=False)
    with open(os.path.join(root, "eval.jsonl"), "w") as f:
        for r in eval_rows:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")

    print(f"\ncorpus: {len(index)} tables -> corpus/records + index.json")
    print(f"eval.jsonl: {len(eval_rows)} queries")
    errs = [r["id"] for r in eval_rows if "error" in r["gold_answer"]]
    if errs:
        print(f"gold-answer errors ({len(errs)}): {errs}")


if __name__ == "__main__":
    main()