#!/usr/bin/env python3
"""One-time build: transform a Common Crawl domain hyperlink graph into a
partitioned Parquet dataset on Hugging Face, queryable remotely for backlinks.

Streams the graph straight from CC's CDN (no 24 GB local copy of the source —
DuckDB reads the gzipped CSVs over httpfs in a single pass). Output (~25-30 GB
of Parquet) is written locally, then pushed to an HF dataset repo.

Run anywhere with ~30 GB free disk + an HF token:
    HF_TOKEN=hf_xxx python3 build.py --repo isam/cc-domain-graph

Produces on the repo:
    manifest.json
    vertices_by_domain.parquet   (id, rev_domain)   sorted by rev_domain
    vertices_by_id.parquet       (id, rev_domain)   sorted by id
    ranks_by_id.parquet          (id, harmonic_pos)
    edges/part=<n>/data_0.parquet (from_id, to_id)  partitioned by to_id // PART_SIZE
"""

import argparse
import json
import os
import shutil
import sys
from pathlib import Path

import duckdb

RELEASE = "cc-main-2025-oct-nov-dec"
BASE = f"https://data.commoncrawl.org/projects/hyperlinkgraph/{RELEASE}/domain"
V_URL = f"{BASE}/{RELEASE}-domain-vertices.txt.gz"
E_URL = f"{BASE}/{RELEASE}-domain-edges.txt.gz"
R_URL = f"{BASE}/{RELEASE}-domain-ranks.txt.gz"

# ~121M domain nodes → ~1M ids per partition → ~122 partition files.
PART_SIZE = 1_000_000


def log(msg: str):
    print(msg, flush=True)


def build(out_dir: Path, memory_limit: str = "12GB") -> dict:
    out_dir.mkdir(parents=True, exist_ok=True)
    (out_dir / "edges").mkdir(exist_ok=True)
    tmp = out_dir / "_duck_tmp"
    tmp.mkdir(exist_ok=True)

    con = duckdb.connect(config={"memory_limit": memory_limit, "temp_directory": str(tmp)})
    con.execute("INSTALL httpfs; LOAD httpfs;")

    # ---- vertices: id, rev_domain (3rd col n_hosts ignored) ----
    log("Loading vertices (851 MiB gz) ...")
    con.execute(f"""
        CREATE TABLE v AS
        SELECT c0::UINTEGER AS id, c1 AS rev_domain
        FROM read_csv('{V_URL}', delim='\t', header=false, auto_detect=false,
                      columns={{'c0':'UINTEGER','c1':'VARCHAR','c2':'UBIGINT'}})
    """)
    n_nodes = con.execute("SELECT count(*) FROM v").fetchone()[0]
    log(f"  {n_nodes:,} domain nodes")
    con.execute(f"COPY (SELECT * FROM v ORDER BY rev_domain) TO '{out_dir}/vertices_by_domain.parquet' (FORMAT parquet)")
    con.execute(f"COPY (SELECT * FROM v ORDER BY id) TO '{out_dir}/vertices_by_id.parquet' (FORMAT parquet)")
    log("  wrote vertices_by_domain.parquet, vertices_by_id.parquet")

    # ---- ranks: harmonicc_pos keyed by host_rev → join to id ----
    log("Loading ranks ...")
    con.execute(f"""
        CREATE TABLE r AS
        SELECT c0::UBIGINT AS harmonic_pos, c4 AS rev_domain
        FROM read_csv('{R_URL}', delim='\t', header=false, skip=1, auto_detect=false,
                      columns={{'c0':'UBIGINT','c1':'VARCHAR','c2':'UBIGINT',
                                'c3':'VARCHAR','c4':'VARCHAR','c5':'UBIGINT'}})
    """)
    con.execute(f"""
        COPY (SELECT v.id, r.harmonic_pos
              FROM r JOIN v USING (rev_domain) ORDER BY v.id)
        TO '{out_dir}/ranks_by_id.parquet' (FORMAT parquet)
    """)
    log("  wrote ranks_by_id.parquet")

    # ---- edges: partition by to_id // PART_SIZE (no global sort needed,
    #      each partition already spans a narrow to_id range) ----
    log("Streaming + partitioning edges (22.9 GiB gz, single pass) — this is the long step ...")
    con.execute(f"""
        COPY (SELECT c0::UINTEGER AS from_id, c1::UINTEGER AS to_id,
                     (c1::UBIGINT // {PART_SIZE})::UINTEGER AS part
              FROM read_csv('{E_URL}', delim='\t', header=false, auto_detect=false,
                            columns={{'c0':'UINTEGER','c1':'UINTEGER'}}))
        TO '{out_dir}/edges' (FORMAT parquet, PARTITION_BY (part), OVERWRITE_OR_IGNORE)
    """)
    n_parts = len(list((out_dir / "edges").glob("part=*")))
    log(f"  wrote {n_parts} edge partitions")

    shutil.rmtree(tmp, ignore_errors=True)

    manifest = {
        "release": RELEASE,
        "source": BASE,
        "part_size": PART_SIZE,
        "domain_nodes": n_nodes,
        "schema": {
            "vertices": ["id", "rev_domain"],
            "edges": ["from_id", "to_id"],
            "ranks": ["id", "harmonic_pos"],
        },
        "note": "rev_domain is label-reversed (blackstump.com.au -> au.com.blackstump). "
                "Backlinks: read edges/part=<to_id//part_size> WHERE to_id=<id>.",
    }
    (out_dir / "manifest.json").write_text(json.dumps(manifest, indent=2))
    log("  wrote manifest.json")
    return manifest


def push(out_dir: Path, repo: str, token: str):
    from huggingface_hub import HfApi
    api = HfApi(token=token)
    api.create_repo(repo, repo_type="dataset", exist_ok=True)
    log(f"Uploading {out_dir} → {repo} (dataset) ...")
    # upload_large_folder handles many files + resumes; fall back to upload_folder.
    try:
        api.upload_large_folder(folder_path=str(out_dir), repo_id=repo, repo_type="dataset")
    except AttributeError:
        api.upload_folder(folder_path=str(out_dir), repo_id=repo, repo_type="dataset")
    log("Done. Query with:  python3 backlinks.py <domain> --repo " + repo)


def main():
    ap = argparse.ArgumentParser(description="Build the hosted CC domain backlink graph.")
    ap.add_argument("--repo", required=True, help="Target HF dataset repo, e.g. isam/cc-domain-graph")
    ap.add_argument("--out", default="./graph_out", help="Local output dir (default: ./graph_out)")
    ap.add_argument("--memory-limit", default="12GB", help="DuckDB memory limit (default: 12GB)")
    ap.add_argument("--no-push", action="store_true", help="Build locally only, skip HF upload")
    args = ap.parse_args()

    out = Path(args.out)
    build(out, memory_limit=args.memory_limit)

    if args.no_push:
        log(f"Built in {out} (not pushed).")
        return
    token = os.environ.get("HF_TOKEN")
    if not token:
        sys.exit("Set HF_TOKEN env var to push (or use --no-push).")
    push(out, args.repo, token)


if __name__ == "__main__":
    main()