#!/usr/bin/env python3
"""WP44 — parse the DBLP XML dump (dblp.xml.gz + dblp.dtd) into Parquet.

Output: <out>/dblp_publications.parquet with one row per publication record
(article, inproceedings, proceedings, book, incollection, phdthesis, mastersthesis,
www is skipped) and <out>/dblp_parse_report.json (counts, md5 of input, timings).
All years are kept (the historical baseline needs 2018–2022; counts per year are
reported for the whole archive).
"""
import argparse, gzip, hashlib, json, os, time, datetime as dt
from lxml import etree
import pyarrow as pa, pyarrow.parquet as pq

TYPES = {"article", "inproceedings", "proceedings", "book", "incollection", "phdthesis", "mastersthesis", "data"}
SCHEMA = pa.schema([
    ("key", pa.string()), ("rec_type", pa.string()), ("mdate", pa.string()), ("publtype", pa.string()),
    ("title", pa.string()), ("year", pa.int32()), ("authors", pa.list_(pa.string())), ("author_pids", pa.list_(pa.string())),
    ("editors", pa.list_(pa.string())), ("journal", pa.string()), ("booktitle", pa.string()), ("volume", pa.string()),
    ("number", pa.string()), ("pages", pa.string()), ("publisher", pa.string()), ("series", pa.string()),
    ("doi", pa.string()), ("ee", pa.list_(pa.string())), ("url", pa.string()), ("crossref", pa.string()),
    ("stream", pa.string()),
])


def text(el):
    return "".join(el.itertext()).strip() if el is not None else None


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--xml", required=True); ap.add_argument("--out", required=True)
    ap.add_argument("--min-year", type=int, default=0)
    a = ap.parse_args(); a.xml = os.path.abspath(a.xml); a.out = os.path.abspath(a.out); os.makedirs(a.out, exist_ok=True)
    t0 = time.time()
    md5 = hashlib.md5(open(a.xml, "rb").read(1 << 30)).hexdigest()  # first GB hash (file is 1.1 GB) + size recorded
    size = os.path.getsize(a.xml)
    # lxml needs the DTD next to the XML for entity resolution; parse the gz stream
    parser = etree.XMLParser(load_dtd=True, resolve_entities=True, huge_tree=True)
    counts = {}; years = {}
    rows = {n: [] for n in SCHEMA.names}; n = 0
    writer = pq.ParquetWriter(os.path.join(a.out, "dblp_publications.parquet"), SCHEMA, compression="zstd")

    def flush():
        nonlocal rows
        if rows["key"]:
            writer.write_table(pa.Table.from_pydict(rows, schema=SCHEMA)); rows = {k: [] for k in SCHEMA.names}

    # iterparse over the gzip stream; the DTD is referenced relatively (dblp.dtd) -> chdir
    os.chdir(os.path.dirname(os.path.abspath(a.xml)))
    with gzip.open(a.xml, "rb") as fh:
        ctx = etree.iterparse(fh, events=("end",), tag=list(TYPES | {"www"}), load_dtd=True, resolve_entities=True, huge_tree=True)
        for _, el in ctx:
            t = el.tag
            if t in TYPES:
                counts[t] = counts.get(t, 0) + 1
                yr = el.findtext("year"); yr = int(yr) if yr and yr.isdigit() else None
                years[yr] = years.get(yr, 0) + 1
                if yr is not None and yr >= a.min_year:
                    ees = [e.text for e in el.findall("ee") if e.text]
                    doi = next((e[16:].lower() for e in ees if e.startswith("https://doi.org/")), None)
                    stream = el.get("key", "").rsplit("/", 1)[0] if "/" in el.get("key", "") else None
                    vals = dict(key=el.get("key"), rec_type=t, mdate=el.get("mdate"), publtype=el.get("publtype"),
                                title=text(el.find("title")), year=yr,
                                authors=[text(x) for x in el.findall("author")], author_pids=[x.get("pid") for x in el.findall("author")],
                                editors=[text(x) for x in el.findall("editor")], journal=el.findtext("journal"), booktitle=el.findtext("booktitle"),
                                volume=el.findtext("volume"), number=el.findtext("number"), pages=el.findtext("pages"),
                                publisher=el.findtext("publisher"), series=text(el.find("series")), doi=doi, ee=ees,
                                url=el.findtext("url"), crossref=el.findtext("crossref"), stream=stream)
                    for k, v in vals.items():
                        rows[k].append(v)
                    n += 1
                    if n % 200000 == 0:
                        flush(); print(f"{n} rows, {time.time()-t0:.0f}s", flush=True)
            el.clear()
            while el.getprevious() is not None:
                del el.getparent()[0]
    flush(); writer.close()
    rep = {"input": os.path.abspath(a.xml), "input_bytes": size, "input_md5_first_GiB": md5, "parsed_at": dt.datetime.now(dt.UTC).isoformat(),
           "rows_written": n, "min_year": a.min_year, "counts_by_type": counts, "counts_by_year": {str(k): v for k, v in sorted(years.items(), key=lambda kv: (kv[0] is None, kv[0]))},
           "seconds": round(time.time() - t0)}
    json.dump(rep, open(os.path.join(a.out, "dblp_parse_report.json"), "w"), indent=1)
    print(json.dumps({k: v for k, v in rep.items() if k != "counts_by_year"}, indent=1))


if __name__ == "__main__":
    main()