#!/usr/bin/env python3
"""Choose which patents to fetch claims for, before spending network calls.

Two modes.

Granted-claim targets (default) — walk a directory of HUPD application JSONs and
emit the patent numbers worth fetching: granted, in a target domain, and with
enough Background/Abstract text to survive the dataset filters. Pre-filtering
here avoids thousands of fetches for applications the builder would discard.

    python3 scripts/select_targets.py --hupd hupd/sample/2016 \\
        --output data/cache/targets_granted.txt

Publication targets (--publications) — emit "<application>\\t<publication>" pairs
for the SFT training split, so the as-filed (A1) claims can be fetched to build
DPO preference pairs. Restricted to the training split so the held-out evaluation
pool is never touched.

    python3 scripts/select_targets.py --hupd hupd/sample/2016 --publications \\
        --manifest data/sft/manifest.jsonl --output data/cache/targets_pubs.tsv
"""

import argparse
import glob
import json
import os
import random
import sys
from collections import Counter

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

MIN_BACKGROUND_CHARS = 400
MIN_ABSTRACT_CHARS = 150


def iter_applications(hupd_dirs):
    for d in hupd_dirs:
        for path in sorted(glob.glob(os.path.join(d, "*.json"))):
            try:
                with open(path) as f:
                    yield json.load(f)
            except ValueError:
                continue


def already_fetched(path, key):
    done = set()
    if path and 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("--hupd", nargs="+", required=True, help="directories of HUPD JSONs")
    ap.add_argument("--output", required=True)
    ap.add_argument("--publications", action="store_true",
                    help="emit application/publication pairs for as-filed claims")
    ap.add_argument("--manifest", help="restrict to the train split of this manifest")
    ap.add_argument("--domains", help="comma-separated domain allowlist")
    ap.add_argument("--skip-cached", help="JSONL of already-fetched records")
    ap.add_argument("--per-domain-cap", type=int, default=0,
                    help="cap per domain, for augmentation runs")
    ap.add_argument("--seed", type=int, default=42)
    args = ap.parse_args()

    allow = set(args.domains.split(",")) if args.domains else None

    restrict = None
    if args.manifest:
        restrict = set()
        with open(args.manifest) as f:
            for line in f:
                r = json.loads(line)
                if r.get("split") == "train":
                    restrict.add(r["application_number"])

    stats = Counter()
    by_domain = {}
    rows = []

    if args.publications:
        done = already_fetched(args.skip_cached, "application_number")
        for d in iter_applications(args.hupd):
            app = d.get("application_number") or ""
            if restrict is not None and app not in restrict:
                continue
            if app in done:
                stats["already_cached"] += 1
                continue
            pub = (d.get("publication_number") or "").split("-")[0]
            if not pub.startswith("US"):
                stats["no_publication"] += 1
                continue
            rows.append("%s\t%s" % (app, pub))
    else:
        done = already_fetched(args.skip_cached, "patent_number")
        for d in iter_applications(args.hupd):
            if d.get("decision") != "ACCEPTED":
                stats["not_granted"] += 1
                continue
            domain = classify(d.get("main_ipcr_label"))
            if not domain or (allow and domain not in allow):
                stats["off_domain"] += 1
                continue
            patent = d.get("patent_number") or ""
            if not patent.isdigit():
                stats["no_patent_number"] += 1
                continue
            if patent in done:
                stats["already_cached"] += 1
                continue
            if len(clean_section(d.get("background"))) < MIN_BACKGROUND_CHARS:
                stats["thin_background"] += 1
                continue
            if len(clean_section(d.get("abstract"))) < MIN_ABSTRACT_CHARS:
                stats["thin_abstract"] += 1
                continue
            by_domain.setdefault(domain, []).append(patent)

        rng = random.Random(args.seed)
        for domain, patents in sorted(by_domain.items()):
            patents = sorted(set(patents))
            rng.shuffle(patents)
            if args.per_domain_cap:
                patents = patents[: args.per_domain_cap]
            print("  %-26s %5d" % (domain, len(patents)))
            rows.extend(patents)

    rows = list(dict.fromkeys(rows))
    os.makedirs(os.path.dirname(os.path.abspath(args.output)), exist_ok=True)
    with open(args.output, "w") as f:
        f.write("\n".join(rows) + ("\n" if rows else ""))

    print("\nskipped: %s" % ("  ".join("%s=%d" % kv for kv in sorted(stats.items())) or "none"))
    print("wrote %d targets -> %s" % (len(rows), args.output))


if __name__ == "__main__":
    main()