"""COLD Cases parquet shard -> one plain-text corpus for the live demo.

Writes data/cold/corpus.txt: one decision per record ("case name, court, date"
header + opinion texts), records separated by blank lines, plus corpus.json.
"""
import json, sys, re, html
import pyarrow.parquet as pq

src, dst = sys.argv[1], sys.argv[2]
limit = int(sys.argv[3]) if len(sys.argv) > 3 else 10**12
pf = pq.ParquetFile(src)
print(pf.schema_arrow)
n_dec = n_bytes = 0
courts = {}
years = []
tag = re.compile(r"<[^>]+>")
with open(dst, "w", encoding="utf-8") as out:
    for rg in range(pf.num_row_groups):
        t = pf.read_row_group(rg, columns=["case_name", "court_short_name", "date_filed", "opinions"]).to_pylist()
        for r in t:
            texts = []
            for o in r["opinions"] or []:
                tx = o.get("opinion_text") or ""
                if "<" in tx and ">" in tx:
                    tx = html.unescape(tag.sub(" ", tx))
                tx = re.sub(r"\n{3,}", "\n\n", tx.strip())
                if tx:
                    texts.append(tx)
            if not texts:
                continue
            head = f"{r['case_name'] or ''}, {r['court_short_name'] or ''}, {r['date_filed'] or ''}"
            rec = head + "\n\n" + "\n\n".join(texts) + "\n\n\n"
            out.write(rec)
            n_dec += 1; n_bytes += len(rec.encode())
            courts[r["court_short_name"]] = courts.get(r["court_short_name"], 0) + 1
            if r["date_filed"]:
                years.append(int(str(r["date_filed"])[:4]))
            if n_bytes >= limit:
                break
        if n_bytes >= limit:
            break
        print(f"row group {rg}: {n_dec:,} decisions, {n_bytes/1e9:.2f} GB", flush=True)
years.sort()
meta = dict(source="harvard-lil/cold-cases (CC0), " + src.split("/")[-1], decisions=n_dec, bytes=n_bytes,
            courts=len(courts), top_courts=sorted(courts.items(), key=lambda kv: -kv[1])[:12],
            years=[years[0], years[len(years)//2], years[-1]] if years else None)
json.dump(meta, open(dst.rsplit(".", 1)[0] + ".json", "w"), indent=1)
print(json.dumps(meta, indent=1))