#!/usr/bin/env python3
"""Run pipeline/sql/*.sql against the NEAR BigQuery public dataset and snapshot results to raw/bigquery/.

Needs GOOGLE_APPLICATION_CREDENTIALS (any Google account works: the sandbox gives 1 TB scanned a month, no card).
Usage: run_bq.py [--only 01,02] [--cap-gb 200]
Without credentials the pipeline still rebuilds from the committed snapshots.
"""
import argparse, csv, datetime, gzip, hashlib, io, json, pathlib, sys

HERE = pathlib.Path(__file__).resolve().parent
ROOT = HERE.parent
sys.path.insert(0, str(HERE))
from bq import BQ  # noqa: E402


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--only", default="")
    ap.add_argument("--cap-gb", type=float, default=200)
    ap.add_argument("--dry", action="store_true")
    ap.add_argument("--total-cap-gb", type=float, default=300, help="stop before this many GB in one run")
    args = ap.parse_args()
    client = BQ(max_bytes=args.cap_gb * 1e9)
    only = [x for x in args.only.split(",") if x]
    outdir = ROOT / "raw" / "bigquery"
    outdir.mkdir(parents=True, exist_ok=True)
    total = 0
    for path in sorted((HERE / "sql").glob("*.sql")):
        if only and path.name.split("_")[0] not in only:
            continue
        sql = path.read_text()
        if "{{INTENTS_TOKENS}}" in sql:
            tokens = [t.strip() for t in (HERE / "sql" / "intents_tokens.txt").read_text().splitlines() if t.strip()]
            sql = sql.replace("{{INTENTS_TOKENS}}", ",".join("'" + t.replace("'", "") + "'" for t in tokens))
        if "{{TOKEN_PRICE_STRUCTS}}" in sql:
            rows = list(csv.DictReader((HERE / "sql" / "token_prices.csv").open()))
            structs = [f"STRUCT('{r['token'].replace(chr(39), '')}' AS token, {int(r['decimals'])} AS decimals, {float(r['price'])} AS price)" for r in rows]
            sql = sql.replace("{{TOKEN_PRICE_STRUCTS}}", ", ".join(structs))
        name = path.stem
        if args.dry:
            gb = client.dry_run(sql) / 1e9
            total += gb
            print(f"{name}: {gb:.2f} GB")
            continue
        if total + client.dry_run(sql) / 1e9 > args.total_cap_gb:
            raise SystemExit(f"stopping before {name}: run budget of {args.total_cap_gb} GB would be passed")
        rows, stats = client.query(sql, name)
        buf = io.StringIO()
        if rows:
            wr = csv.DictWriter(buf, fieldnames=list(rows[0].keys()), lineterminator="\n")
            wr.writeheader()
            wr.writerows(rows)
        raw = buf.getvalue().encode()
        dest = outdir / f"{name}.csv.gz"
        with gzip.GzipFile(dest, "wb", mtime=0) as fh:
            fh.write(raw)
        meta = {"file": dest.name, "pulled_at": datetime.datetime.now(datetime.timezone.utc).strftime("%Y-%m-%dT%H:%MZ"),
                "rows": len(rows), "sha256_csv": hashlib.sha256(raw).hexdigest(),
                "sql_sha256": hashlib.sha256(sql.encode()).hexdigest(), "bytes_scanned": stats["bytes_processed"]}
        with (outdir / "MANIFEST.jsonl").open("a", encoding="utf8") as fh:
            fh.write(json.dumps(meta) + "\n")
        total += stats["bytes_processed"] / 1e9
        print(f"{name}: {len(rows)} rows, {stats['bytes_processed']/1e9:.2f} GB scanned, {dest.stat().st_size/1024:.0f} KB")
    print(f"total {total:.2f} GB")


if __name__ == "__main__":
    main()