#!/usr/bin/env python3
"""WP44 — OpenAlex public snapshot ingestion (works entity).

Streams every part file of the works snapshot listed in the manifest, verifies
content length and records a SHA-256 of the gzip bytes (immutable provenance),
parses each JSON line and writes:

  * ``cs/<partition>__<part>.jsonl.zst``  — the FULL OpenAlex record of every
    work matched by the CS filter (any topic in field 17 Computer Science,
    publication_year in [YEAR_MIN, YEAR_MAX]); XPAC works are kept with the
    ``is_xpac`` flag so that the corpus filter can be applied later.
  * ``slim/<partition>__<part>.parquet``  — a slim row for EVERY work in the
    snapshot (identifiers, year, type, primary topic, counts) used to resolve
    cited works outside computer science.

A SQLite ledger makes the job resumable and keeps job-level provenance
(attempts, timings, hashes, counts, errors). Raw .gz files are deleted after a
successful parse unless --keep-raw is given (the snapshot is a public,
versioned release; the manifest and per-file hashes are kept).
"""
from __future__ import annotations

import argparse
import datetime as dt
import gzip
import hashlib
import json
import multiprocessing as mp
import os
import sqlite3
import subprocess
import sys
import time
import traceback

import orjson
import pyarrow as pa
import pyarrow.parquet as pq
import zstandard as zstd

S3_HTTP = "https://openalex.s3.amazonaws.com/"
ID_SET = None  # optional: set of OpenAlex work ids (ints) to extract instead of the CS filter (second pass)
CS_FIELD = "https://openalex.org/fields/17"
YEAR_MIN, YEAR_MAX = 2018, 2026

SLIM_SCHEMA = pa.schema([
    ("id", pa.int64()),
    ("doi", pa.string()),
    ("publication_year", pa.int32()),
    ("publication_date", pa.string()),
    ("type", pa.string()),
    ("cited_by_count", pa.int32()),
    ("primary_topic_id", pa.int32()),
    ("primary_subfield_id", pa.int32()),
    ("primary_field_id", pa.int32()),
    ("primary_domain_id", pa.int32()),
    ("cs_any_topic", pa.bool_()),
    ("is_xpac", pa.bool_()),
    ("is_oa", pa.bool_()),
    ("is_retracted", pa.bool_()),
    ("language", pa.string()),
    ("source_id", pa.int64()),
    ("n_refs", pa.int32()),
    ("n_authors", pa.int32()),
    ("has_abstract", pa.bool_()),
    ("fwci", pa.float32()),
])


def _oa_int(url, prefix_len=None):
    """'https://openalex.org/W123' -> 123 (None-safe)."""
    if not url:
        return None
    tail = url.rsplit("/", 1)[-1]
    # strip leading letter(s): W, T, I, S...
    i = 0
    while i < len(tail) and not tail[i].isdigit():
        i += 1
    try:
        return int(tail[i:])
    except ValueError:
        return None


def slim_row(w: dict) -> tuple:
    pt = w.get("primary_topic") or {}
    topics = w.get("topics") or []
    cs_any = any(((t.get("field") or {}).get("id") == CS_FIELD) for t in topics) or (
        (pt.get("field") or {}).get("id") == CS_FIELD)
    loc = w.get("primary_location") or {}
    src = (loc.get("source") or {}) if loc else {}
    oa = w.get("open_access") or {}
    doi = w.get("doi")
    if doi and doi.startswith("https://doi.org/"):
        doi = doi[16:]
    return (
        _oa_int(w.get("id")),
        doi.lower() if doi else None,
        w.get("publication_year"),
        w.get("publication_date"),
        w.get("type"),
        w.get("cited_by_count"),
        _oa_int(pt.get("id")),
        _oa_int((pt.get("subfield") or {}).get("id")),
        _oa_int((pt.get("field") or {}).get("id")),
        _oa_int((pt.get("domain") or {}).get("id")),
        bool(cs_any),
        bool(w.get("is_xpac")),
        bool(oa.get("is_oa")),
        bool(w.get("is_retracted")),
        w.get("language"),
        _oa_int(src.get("id")) if src else None,
        len(w.get("referenced_works") or []),
        w.get("authors_count") if w.get("authors_count") is not None else len(w.get("authorships") or []),
        w.get("abstract_inverted_index") is not None,
        w.get("fwci"),
    ), cs_any


def _load_ids(path):
    with open(path) as fh:
        next(fh)
        return {int(x.strip().rsplit("W", 1)[-1]) for x in fh if x.strip()}


def _init_worker(id_list):
    """Pool initializer: macOS uses the spawn start method, so module globals must be rebuilt in each worker."""
    global ID_SET
    ID_SET = _load_ids(id_list) if id_list else None


def process_file(args):
    """Download, hash, parse and filter one part file. Returns a result dict."""
    url, expected_len, out_dir, keep_raw, mode = args
    rel = url.replace("s3://openalex/data/jsonl/works/", "")
    name = rel.replace("/", "__")                      # updated_date=X__part_0000.gz
    http = url.replace("s3://openalex/", S3_HTTP)
    tmp_dir = os.path.join(out_dir, "tmp")
    cs_dir = os.path.join(out_dir, "cs")
    slim_dir = os.path.join(out_dir, "slim")
    for d in (tmp_dir, cs_dir, slim_dir):
        os.makedirs(d, exist_ok=True)
    raw_path = os.path.join(tmp_dir, name)
    t0 = time.time()
    res = {"url": url, "name": name, "expected_len": expected_len}
    try:
        # ---- download (curl, retries, resume) ----
        for attempt in range(1, 6):
            cmd = ["curl", "-sS", "-L", "--retry", "5", "--retry-delay", "3", "--retry-all-errors",
                   "-C", "-", "-o", raw_path, http]
            p = subprocess.run(cmd, capture_output=True, text=True, timeout=3600)
            if os.path.exists(raw_path) and os.path.getsize(raw_path) == expected_len:
                break
            time.sleep(5 * attempt)
        got = os.path.getsize(raw_path) if os.path.exists(raw_path) else -1
        if got != expected_len:
            raise RuntimeError(f"size mismatch: got {got}, expected {expected_len}")
        t_dl = time.time() - t0
        # ---- hash ----
        h = hashlib.sha256()
        with open(raw_path, "rb") as fh:
            for chunk in iter(lambda: fh.read(1 << 22), b""):
                h.update(chunk)
        sha = h.hexdigest()
        # ---- parse ----
        cols = [[] for _ in SLIM_SCHEMA]
        n_all = n_cs = n_bad = 0
        cs_tmp = os.path.join(cs_dir, name[:-3] + ".jsonl.zst.part")
        cctx = zstd.ZstdCompressor(level=6)
        with open(cs_tmp, "wb") as cs_fh, cctx.stream_writer(cs_fh) as cs_w, gzip.open(raw_path, "rb") as g:
            for line in g:
                if not line.strip():
                    continue
                try:
                    w = orjson.loads(line)
                except Exception:
                    n_bad += 1
                    continue
                n_all += 1
                if mode == "idset":
                    if _oa_int(w.get("id")) in ID_SET:
                        cs_w.write(line if line.endswith(b"\n") else line + b"\n")
                        n_cs += 1
                    continue
                row, cs_any = slim_row(w)
                for i, v in enumerate(row):
                    cols[i].append(v)
                y = w.get("publication_year")
                if cs_any and y is not None and YEAR_MIN <= y <= YEAR_MAX:
                    cs_w.write(line if line.endswith(b"\n") else line + b"\n")
                    n_cs += 1
        os.replace(cs_tmp, cs_tmp[:-5])
        if mode == "idset":
            cols = None
        table = None if cols is None else pa.Table.from_arrays([pa.array(c, type=f.type) for c, f in zip(cols, SLIM_SCHEMA)], schema=SLIM_SCHEMA)
        if table is not None:
            slim_tmp = os.path.join(slim_dir, name[:-3] + ".parquet.part")
            pq.write_table(table, slim_tmp, compression="zstd")
            os.replace(slim_tmp, slim_tmp[:-5])
        if not keep_raw:
            os.remove(raw_path)
        res.update(status="done", sha256=sha, n_records=n_all, n_cs=n_cs, n_bad=n_bad,
                   t_download=round(t_dl, 1), t_total=round(time.time() - t0, 1), error=None)
    except Exception as e:  # noqa: BLE001
        res.update(status="error", error=f"{e}\n{traceback.format_exc()[-800:]}", t_total=round(time.time() - t0, 1))
    return res


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out", required=True)
    ap.add_argument("--manifest", required=True, help="local copy of works manifest.json")
    ap.add_argument("--workers", type=int, default=6)
    ap.add_argument("--keep-raw", action="store_true")
    ap.add_argument("--limit", type=int, default=0, help="process only N files (smoke test)")
    ap.add_argument("--smallest-first", action="store_true")
    ap.add_argument("--id-list", default=None, help="CSV with a header and one OpenAlex work id per line: extract exactly these works (second pass)")
    a = ap.parse_args()
    global ID_SET
    mode = "cs"
    if a.id_list:
        ID_SET = _load_ids(a.id_list)
        mode = "idset"
        print(f"id-set mode: {len(ID_SET)} ids", flush=True)
    os.makedirs(a.out, exist_ok=True)
    man = json.load(open(a.manifest))
    files = [(f["url"], f["meta"]["content_length"], f["meta"].get("record_count")) for f in man["files"]]
    if a.smallest_first:
        files.sort(key=lambda x: x[1])
    if a.limit:
        files = files[: a.limit]
    db = sqlite3.connect(os.path.join(a.out, "ledger.sqlite"))
    db.execute("""CREATE TABLE IF NOT EXISTS files(url TEXT PRIMARY KEY, name TEXT, expected_len INTEGER,
        manifest_records INTEGER, status TEXT, sha256 TEXT, n_records INTEGER, n_cs INTEGER, n_bad INTEGER,
        t_download REAL, t_total REAL, attempts INTEGER DEFAULT 0, started TEXT, finished TEXT, error TEXT)""")
    db.execute("""CREATE TABLE IF NOT EXISTS runs(started TEXT, manifest_date TEXT, manifest_records INTEGER,
        manifest_bytes INTEGER, workers INTEGER, host TEXT, argv TEXT)""")
    db.execute("INSERT INTO runs VALUES(?,?,?,?,?,?,?)", (dt.datetime.now(dt.UTC).isoformat(), man.get("date"),
               man.get("record_count"), man.get("content_length"), a.workers, os.uname().nodename, " ".join(sys.argv)))
    for url, ln, rc in files:
        db.execute("INSERT OR IGNORE INTO files(url,name,expected_len,manifest_records,status) VALUES(?,?,?,?,'pending')",
                   (url, url.replace("s3://openalex/data/jsonl/works/", "").replace("/", "__"), ln, rc))
    db.commit()
    done = {r[0] for r in db.execute("SELECT url FROM files WHERE status='done'")}
    todo = [(u, ln, a.out, a.keep_raw, mode) for u, ln, _ in files if u not in done]
    total_bytes = sum(t[1] for t in todo)
    print(f"[{dt.datetime.now():%H:%M:%S}] manifest {man.get('date')}: {len(files)} files, {len(done)} done, "
          f"{len(todo)} to do ({total_bytes/1e9:.1f} GB), workers={a.workers}", flush=True)
    t0 = time.time(); done_bytes = 0; n_ok = n_err = 0
    with mp.Pool(a.workers, maxtasksperchild=50, initializer=_init_worker, initargs=(a.id_list,)) as pool:
        for r in pool.imap_unordered(process_file, todo):
            now = dt.datetime.now(dt.UTC).isoformat()
            db.execute("UPDATE files SET attempts=attempts+1 WHERE url=?", (r["url"],))
            if r["status"] == "done":
                n_ok += 1; done_bytes += r["expected_len"]
                db.execute("""UPDATE files SET status='done', sha256=?, n_records=?, n_cs=?, n_bad=?, t_download=?,
                    t_total=?, finished=?, error=NULL WHERE url=?""",
                           (r["sha256"], r["n_records"], r["n_cs"], r["n_bad"], r["t_download"], r["t_total"], now, r["url"]))
            else:
                n_err += 1
                db.execute("UPDATE files SET status='error', error=?, finished=? WHERE url=?", (r["error"], now, r["url"]))
            db.commit()
            el = time.time() - t0
            rate = done_bytes / el if el else 0
            eta = (total_bytes - done_bytes) / rate / 60 if rate else float("nan")
            print(f"[{dt.datetime.now():%H:%M:%S}] {r['status']:5s} {r['name']:40s} rec={r.get('n_records')} "
                  f"cs={r.get('n_cs')} | {n_ok} ok {n_err} err | {done_bytes/1e9:.1f}/{total_bytes/1e9:.1f} GB "
                  f"{rate/1e6:.1f} MB/s ETA {eta:.0f} min", flush=True)
    print("finished", n_ok, "ok", n_err, "errors", flush=True)


if __name__ == "__main__":
    main()