#!/usr/bin/env python3
"""Stream an additional HUPD year archive and keep only under-represented domains.

The Jan-2016 sample yields very few biotech / energy / pharma applications. This
streams a full-year archive straight from Hugging Face without ever writing the
multi-GB tarball to disk (disk here is tight), decompressing member-by-member and
saving only ACCEPTED applications in the requested domains.

Usage:
    python3 augment_thin_domains.py <year> <out_dir> [--domains d1,d2] [--max N]
"""

import argparse
import io
import json
import os
import sys
import tarfile
import time
import urllib.request

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from claim_drafter.domains import classify

URL = "https://huggingface.co/datasets/HUPD/hupd/resolve/main/data/%s.tar.gz"
DEFAULT_DOMAINS = "biotech_life_sciences,energy_environment,pharma_medical_devices"


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("year")
    ap.add_argument("out_dir")
    ap.add_argument("--domains", default=DEFAULT_DOMAINS)
    ap.add_argument("--max", type=int, default=0)
    args = ap.parse_args()

    wanted = set(args.domains.split(","))
    os.makedirs(args.out_dir, exist_ok=True)

    req = urllib.request.Request(URL % args.year, headers={"User-Agent": "Mozilla/5.0"})
    resp = urllib.request.urlopen(req, timeout=120)
    # 'r|gz' = streaming mode: sequential, no seeking, constant memory.
    tar = tarfile.open(fileobj=resp, mode="r|gz")

    seen = kept = 0
    by_domain = {}
    t0 = time.time()
    for member in tar:
        if not member.name.endswith(".json"):
            continue
        seen += 1
        f = tar.extractfile(member)
        if f is None:
            continue
        try:
            d = json.loads(f.read().decode("utf8", "ignore"))
        except ValueError:
            continue
        if d.get("decision") != "ACCEPTED":
            continue
        dom = classify(d.get("main_ipcr_label"))
        if dom not in wanted:
            continue
        if not (d.get("patent_number") or "").isdigit():
            continue
        app = d.get("application_number") or str(kept)
        with open(os.path.join(args.out_dir, app + ".json"), "w") as out:
            json.dump(d, out)
        kept += 1
        by_domain[dom] = by_domain.get(dom, 0) + 1
        if kept % 200 == 0:
            el = time.time() - t0
            print("  scanned %d, kept %d  (%.0fs, %.0f apps/s)  %s"
                  % (seen, kept, el, seen / max(el, 1e-9), by_domain), flush=True)
        if args.max and kept >= args.max:
            break

    print("Done: scanned %d, kept %d -> %s" % (seen, kept, args.out_dir))
    print("By domain: %s" % by_domain)


if __name__ == "__main__":
    main()