#!/usr/bin/env python3
"""Fetch claim text from Google Patents, concurrently and resumably.

Two modes:

    --kind granted   input: one patent number per line
                     output: the claims an examiner ALLOWED (B2/B1 publication)

    --kind filed     input: "<application_number>\\t<publication_number>" per line
                     output: the claims as originally FILED (A1 publication)

Both write JSONL and skip anything already present in the output file, so an
interrupted run can simply be re-run.

    python3 scripts/fetch_claims.py --kind granted \\
        --input data/cache/targets_granted.txt \\
        --output data/cache/granted_claims.jsonl

Please keep --workers modest. This is a polite crawler with backoff, not a
scraper; the underlying patent text is US government public domain, but the
service serving it is not yours.
"""

import argparse
import json
import os
import queue
import random
import sys
import threading
import time

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from claim_drafter.patents import fetch_filed_claims, fetch_granted_claims


def load_done(path, key):
    done = set()
    if os.path.exists(path):
        with open(path) as f:
            for line in f:
                try:
                    done.add(json.loads(line)[key])
                except Exception:
                    pass
    return done


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--kind", choices=["granted", "filed"], required=True)
    ap.add_argument("--input", required=True)
    ap.add_argument("--output", required=True)
    ap.add_argument("--workers", type=int, default=8)
    ap.add_argument("--limit", type=int, default=0)
    args = ap.parse_args()

    os.makedirs(os.path.dirname(os.path.abspath(args.output)), exist_ok=True)
    id_key = "patent_number" if args.kind == "granted" else "application_number"
    done = load_done(args.output, id_key)

    todo = []
    with open(args.input) as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            if args.kind == "granted":
                if line.isdigit() and line not in done:
                    todo.append((line,))
            else:
                parts = line.split("\t")
                if len(parts) == 2 and parts[0] not in done:
                    todo.append(tuple(parts))
    todo = list(dict.fromkeys(todo))
    if args.limit:
        todo = todo[: args.limit]

    print("%d already cached, %d to fetch, %d workers"
          % (len(done), len(todo), args.workers), flush=True)
    if not todo:
        return

    q = queue.Queue()
    for item in todo:
        q.put(item)
    lock = threading.Lock()
    out = open(args.output, "a")
    counts = {"ok": 0, "fail": 0}
    t0 = time.time()

    def worker():
        while True:
            try:
                item = q.get_nowait()
            except queue.Empty:
                return
            if args.kind == "granted":
                record = fetch_granted_claims(item[0])
            else:
                claims = fetch_filed_claims(item[1])
                record = ({"application_number": item[0],
                           "publication_number": item[1],
                           "claims": claims} if claims else None)
            with lock:
                if record:
                    out.write(json.dumps(record) + "\n")
                    out.flush()
                    counts["ok"] += 1
                else:
                    counts["fail"] += 1
                n = counts["ok"] + counts["fail"]
                if n % 100 == 0 or n == len(todo):
                    rate = n / max(time.time() - t0, 1e-9)
                    print("  %d/%d  ok=%d fail=%d  %.1f/s  eta %.1fm"
                          % (n, len(todo), counts["ok"], counts["fail"], rate,
                             (len(todo) - n) / max(rate, 1e-9) / 60), flush=True)
            time.sleep(0.15 + random.random() * 0.25)

    threads = [threading.Thread(target=worker, daemon=True) for _ in range(args.workers)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()
    out.close()
    print("Done: %d fetched, %d failed -> %s" % (counts["ok"], counts["fail"], args.output))


if __name__ == "__main__":
    main()