#!/usr/bin/env python3
"""Download the small OpenAlex entity snapshots (topics, subfields, fields, domains,
sources, institutions, publishers, funders, concepts) as JSONL and convert them to
parquet-friendly JSONL.zst files with a SHA-256 ledger. Authors are NOT downloaded
(embedded authorship data is used); keywords are optional (--with-keywords)."""
import argparse, hashlib, json, os, subprocess, time, datetime as dt

S3_HTTP = "https://openalex.s3.amazonaws.com/"
ENTITIES = ["topics", "subfields", "fields", "domains", "sources", "institutions", "publishers", "funders", "concepts"]


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--out", required=True); ap.add_argument("--with-keywords", action="store_true")
    a = ap.parse_args()
    ents = ENTITIES + (["keywords"] if a.with_keywords else [])
    ledger_path = os.path.join(a.out, "entities_ledger.json")
    ledger = json.load(open(ledger_path)) if os.path.exists(ledger_path) else {}
    for e in ents:
        man = json.load(open(os.path.join(a.out, f"{e}_manifest.json")))
        d = os.path.join(a.out, e); os.makedirs(d, exist_ok=True)
        for f in man["files"]:
            url = f["url"]; ln = f["meta"]["content_length"]
            name = url.split(f"/{e}/")[1].replace("/", "__")
            path = os.path.join(d, name)
            if ledger.get(url, {}).get("status") == "done" and os.path.exists(path):
                continue
            for attempt in range(5):
                subprocess.run(["curl", "-sS", "-L", "--retry", "5", "-C", "-", "-o", path, url.replace("s3://openalex/", S3_HTTP)])
                if os.path.exists(path) and os.path.getsize(path) == ln:
                    break
                time.sleep(3)
            ok = os.path.exists(path) and os.path.getsize(path) == ln
            sha = hashlib.sha256(open(path, "rb").read()).hexdigest() if ok else None
            ledger[url] = {"entity": e, "path": path, "bytes": ln, "sha256": sha, "status": "done" if ok else "error",
                           "records": f["meta"].get("record_count"), "when": dt.datetime.now(dt.UTC).isoformat()}
            json.dump(ledger, open(ledger_path, "w"), indent=1)
            print(e, name, "ok" if ok else "ERROR", flush=True)


if __name__ == "__main__":
    main()