#!/usr/bin/env python3
"""WP44 — parse the ACL Anthology BibTeX export (anthology+abstracts.bib.gz) into Parquet."""
import argparse, gzip, hashlib, json, os, re, datetime as dt
import pyarrow as pa, pyarrow.parquet as pq

FIELDS = ["title", "author", "editor", "booktitle", "journal", "month", "year", "address", "publisher", "url", "doi", "pages", "abstract", "volume", "number"]
SCHEMA = pa.schema([("acl_id", pa.string()), ("entry_type", pa.string())] + [(f, pa.string()) for f in FIELDS])


def entries(txt):
    i = 0; n = len(txt)
    while True:
        j = txt.find("@", i)
        if j < 0:
            return
        k = txt.find("{", j)
        et = txt[j + 1:k].strip().lower()
        depth = 0; p = k
        while p < n:
            c = txt[p]
            if c == "{":
                depth += 1
            elif c == "}":
                depth -= 1
                if depth == 0:
                    break
            p += 1
        body = txt[k + 1:p]; i = p + 1
        if et in ("comment", "string", "preamble"):
            continue
        yield et, body


def parse_body(body):
    key, _, rest = body.partition(",")
    out = {"acl_id": key.strip()}
    # fields of the form name = {value} or name = "value", values may contain nested braces
    pos = 0
    while True:
        m = re.search(r"(\w+)\s*=\s*", rest[pos:])
        if not m:
            break
        name = m.group(1).lower(); p = pos + m.end()
        if p >= len(rest):
            break
        if rest[p] == "{":
            depth = 0; q = p
            while q < len(rest):
                if rest[q] == "{": depth += 1
                elif rest[q] == "}":
                    depth -= 1
                    if depth == 0: break
                q += 1
            val = rest[p + 1:q]; pos = q + 1
        elif rest[p] == '"':
            q = rest.find('"', p + 1); val = rest[p + 1:q]; pos = q + 1
        else:
            q = rest.find(",", p); q = len(rest) if q < 0 else q; val = rest[p:q]; pos = q
        val = re.sub(r"\s+", " ", val).strip()
        val = val.replace("{", "").replace("}", "")
        if name in FIELDS:
            out[name] = val
    return out


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--bib", required=True); ap.add_argument("--out", required=True)
    a = ap.parse_args(); os.makedirs(a.out, exist_ok=True)
    raw = open(a.bib, "rb").read(); sha = hashlib.sha256(raw).hexdigest()
    txt = gzip.decompress(raw).decode("utf-8", "replace")
    rows = {n: [] for n in SCHEMA.names}; n = 0; types = {}
    for et, body in entries(txt):
        d = parse_body(body); d["entry_type"] = et; types[et] = types.get(et, 0) + 1
        for k in SCHEMA.names:
            rows[k].append(d.get(k))
        n += 1
    pq.write_table(pa.Table.from_pydict(rows, schema=SCHEMA), os.path.join(a.out, "acl_anthology.parquet"), compression="zstd")
    rep = {"input": a.bib, "sha256": sha, "bytes": len(raw), "entries": n, "types": types, "parsed_at": dt.datetime.now(dt.UTC).isoformat()}
    json.dump(rep, open(os.path.join(a.out, "acl_parse_report.json"), "w"), indent=1); print(json.dumps(rep, indent=1))


if __name__ == "__main__":
    main()