"""Stage and upload the splits + the SQLite databases they use to a private HF dataset repo.

Rows are rewritten so db_path is RELATIVE to the dataset root
(databases/<db_id>/<db_id>.sqlite): the same rows then work on the laptop and inside an
HF Job that mounts the repo.

Usage: python build_hf_dataset.py <repo_id> [--no-upload]
Token is read from the HF_TOKEN environment variable (never from a file).
"""

from __future__ import annotations

import json
import os
import shutil
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.abspath(os.path.join(HERE, "..", ".."))
SPLITS = os.path.join(ROOT, "data", "splits")
STAGE = os.path.join(ROOT, "data", "hf_dataset")

CARD = """---
license: cc-by-sa-4.0
task_categories:
- text-generation
language:
- en
tags:
- text-to-sql
- spider
- self-healing
pretty_name: Self-heal SQL - Spider splits
---

# Self-heal SQL: database-disjoint Spider splits

Splits and SQLite databases used for a self-healing text-to-SQL agent
(detect failures -> classify -> generate corrective examples -> DPO retrain -> gate).

Derived from **Spider** (Yu et al., 2018, https://yale-lily.github.io/spider), CC BY-SA 4.0.

| split | source | purpose |
|---|---|---|
| `pool` | Spider train, DB group A | failure pool; its failures become training data |
| `gate` | Spider train, DB group B | held-out deployment gate; never used to build training data |
| `test` | Spider dev | headline before/after numbers |

No database appears in more than one split (asserted at build time). Rows whose gold query
errors, times out or returns zero rows were dropped.

Row fields: `qid, split, db_id, question, gold_sql, db_path, schema`.
`db_path` is relative to the repo root.
"""


def main() -> None:
    repo_id = sys.argv[1]
    upload = "--no-upload" not in sys.argv

    if os.path.isdir(STAGE):
        shutil.rmtree(STAGE)
    os.makedirs(os.path.join(STAGE, "splits"))

    used: dict[str, str] = {}
    for name in ("pool", "gate", "test"):
        rows = []
        with open(os.path.join(SPLITS, f"{name}.jsonl"), encoding="utf-8") as f:
            for line in f:
                r = json.loads(line)
                used[r["db_id"]] = r["db_path"]
                r["db_path"] = f"databases/{r['db_id']}/{r['db_id']}.sqlite"
                rows.append(r)
        with open(os.path.join(STAGE, "splits", f"{name}.jsonl"), "w", encoding="utf-8") as f:
            for r in rows:
                f.write(json.dumps(r, ensure_ascii=False) + "\n")
        print(f"{name}: {len(rows)} rows")

    total = 0
    for db_id, src in used.items():
        dst_dir = os.path.join(STAGE, "databases", db_id)
        os.makedirs(dst_dir)
        dst = os.path.join(dst_dir, f"{db_id}.sqlite")
        shutil.copyfile(src, dst)
        total += os.path.getsize(dst)
    print(f"{len(used)} databases, {total / 2**20:.1f} MiB")

    with open(os.path.join(STAGE, "README.md"), "w", encoding="utf-8") as f:
        f.write(CARD)

    if not upload:
        print("staged at", STAGE, "(not uploaded)")
        return

    from huggingface_hub import HfApi

    api = HfApi(token=os.environ["HF_TOKEN"])
    api.create_repo(repo_id, repo_type="dataset", private=True, exist_ok=True)
    api.upload_folder(repo_id=repo_id, repo_type="dataset", folder_path=STAGE,
                      commit_message="Spider-derived DB-disjoint splits + databases")
    print("uploaded: https://huggingface.co/datasets/" + repo_id, "(private)")


if __name__ == "__main__":
    main()