#!/usr/bin/env python3
"""
nl-handles-scan.py -- enumerate atproto accounts by the country of their CURRENT handle.

One pass over the PLC directory export buckets every European ccTLD at once. The export is a
single global stream and the filter is client-side, so scanning for 50 countries costs the same
as scanning for one: the expensive part is the stream, not the matching. Never run this per
country.

A ccTLD handle is a CONFIRMATION signal -- unlike language, which only ever feeds the estimate
layer. Choosing `.nl` for your handle is a statement; posting in Dutch is not. That is why the
suffix map (scripts/country-tlds.json) is deliberately conservative and excludes word-play TLDs.

Writes, per run:
  <out-prefix>.dids        DIDs for the legacy country (--legacy-country, default nl).
                           Unchanged in name and content -- the NL aggregator reads this exact
                           file, so it keeps working untouched.
  <out-prefix>.<cc>.dids   DIDs for every country seen, ISO 3166-1 alpha-2 lowercase.

Source of truth: https://plc.directory/export, streamed in createdAt order. For each DID we track
its latest handle. Changing to a mapped TLD adds the DID to that country; changing away removes
it; moving from one country to another moves it; a tombstone drops it. Only matched DIDs are held
in memory (a dict did -> country, ~millions at most), never the ~40M-DID universe.

Resumable: the last-seen createdAt is persisted to <out-prefix>.cursor, and every per-country
.dids file is reloaded on start, so re-runs continue incrementally instead of rescanning. The
first full pass takes hours; later runs only process new operations. Results are flushed
periodically, so aggregators see growing counts while the first pass is still running.

PARTIAL RESULTS OVERCOUNT: a DID that later changed away from a mapped TLD is not removed until
the scan reaches that operation. Counts are only trustworthy after a complete pass.

Op handling:
  - plc_operation : handle = first `at://…` entry in operation.alsoKnownAs
  - legacy create : operation.handle
  - plc_tombstone : drop the DID
  - nullified     : skipped

Usage:
  python3 nl-handles-scan.py --out-prefix nl-handles [--sleep 0.15] [--max-pages N]
  python3 nl-handles-scan.py --countries nl,be,de      # restrict the buckets written
  python3 nl-handles-scan.py --tlds .nl,.amsterdam,.frl  # legacy single-country mode
"""
import argparse
import json
import os
import time
import urllib.error
import urllib.request

EXPORT = "https://plc.directory/export"
DEFAULT_TLD_MAP = os.path.join(os.path.dirname(os.path.abspath(__file__)), "country-tlds.json")


def load_tld_map(path):
    """suffix label -> ISO 3166-1 alpha-2. Matching is on the handle's final label only, so
    every key is a single label ('nl', 'amsterdam'), never a dotted suffix."""
    with open(path) as f:
        data = json.load(f)
    return {k.lower().lstrip("."): v.lower() for k, v in data["tlds"].items()}


def handle_of(rec):
    """Return (handle_lowercase_or_None, is_tombstone) for a PLC export record."""
    op = rec.get("operation") or {}
    if op.get("type") == "plc_tombstone":
        return None, True
    aka = op.get("alsoKnownAs")
    if isinstance(aka, list):
        for a in aka:
            if isinstance(a, str) and a.startswith("at://"):
                return a[len("at://"):].lower(), False
    h = op.get("handle")  # legacy `create` ops
    if isinstance(h, str):
        return h.lower().lstrip("@"), False
    return None, False


def fetch_page(after, count=1000, retries=5):
    url = f"{EXPORT}?count={count}"
    if after:
        url += f"&after={after}"
    for i in range(retries):
        try:
            req = urllib.request.Request(
                url, headers={"accept": "application/jsonl", "user-agent": "atproto.nl-nl-handles/1.0"}
            )
            with urllib.request.urlopen(req, timeout=60) as r:
                return r.read().decode("utf-8", "replace")
        except (urllib.error.URLError, TimeoutError, ConnectionError):
            time.sleep(2 ** i)
    raise RuntimeError(f"PLC export fetch failed after {retries} tries: {url}")


def country_of(handle, tld_map):
    """Country for a handle, by its final label. None when the TLD is not mapped."""
    if not handle or "." not in handle:
        return None
    return tld_map.get(handle.rsplit(".", 1)[-1])


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--out-prefix", default="nl-handles")
    ap.add_argument("--dir", default=".")
    ap.add_argument("--tld-map", default=DEFAULT_TLD_MAP,
                    help="JSON suffix->country map (default: scripts/country-tlds.json)")
    ap.add_argument("--countries", default="",
                    help="comma-separated alpha-2 allowlist; empty = every country in the map")
    ap.add_argument("--legacy-country", default="nl",
                    help="country written to the unsuffixed <out-prefix>.dids the NL aggregator reads")
    # Legacy single-country mode: an explicit suffix list overrides the map entirely.
    ap.add_argument("--tlds", default="", help="comma-separated handle suffixes (legacy mode)")
    ap.add_argument("--sleep", type=float, default=0.15)
    ap.add_argument("--max-pages", type=int, default=0, help="0 = run to the end of the export")
    ap.add_argument("--flush-every", type=int, default=200)
    args = ap.parse_args()

    if args.tlds:
        tld_map = {t.strip().lower().lstrip("."): args.legacy_country
                   for t in args.tlds.split(",") if t.strip()}
    else:
        tld_map = load_tld_map(args.tld_map)
    if args.countries:
        keep = {c.strip().lower() for c in args.countries.split(",") if c.strip()}
        tld_map = {k: v for k, v in tld_map.items() if v in keep}
    if not tld_map:
        raise SystemExit("no TLDs to match after filtering")

    cursor_path = os.path.join(args.dir, f"{args.out_prefix}.cursor")
    legacy_path = os.path.join(args.dir, f"{args.out_prefix}.dids")

    def country_path(cc):
        return os.path.join(args.dir, f"{args.out_prefix}.{cc}.dids")

    # did -> country, for matched DIDs only. One dict rather than a set per country: a handle
    # change then costs one pop instead of a discard against every country in the map.
    matched = {}

    def load(path, cc):
        if not os.path.exists(path):
            return
        with open(path) as f:
            for ln in f:
                did = ln.strip()
                if did.startswith("did:"):
                    matched[did] = cc

    for cc in sorted(set(tld_map.values())):
        load(country_path(cc), cc)
    # The legacy file predates the per-country ones; read it too so an existing NL scan resumes
    # instead of starting from zero.
    load(legacy_path, args.legacy_country)

    after = None
    if os.path.exists(cursor_path):
        after = open(cursor_path).read().strip() or None

    def write_dids(path, dids):
        tmp = path + ".tmp"
        with open(tmp, "w") as f:
            f.write("\n".join(sorted(dids)))
            if dids:
                f.write("\n")
        os.replace(tmp, path)

    def flush():
        by_country = {}
        for did, cc in matched.items():
            by_country.setdefault(cc, []).append(did)
        for cc in sorted(set(tld_map.values())):
            write_dids(country_path(cc), by_country.get(cc, []))
        write_dids(legacy_path, by_country.get(args.legacy_country, []))
        if after:
            with open(cursor_path, "w") as f:
                f.write(after)
        return by_country

    def summary(by_country, top=6):
        counts = sorted(((cc, len(v)) for cc, v in by_country.items()), key=lambda kv: -kv[1])
        return " ".join(f"{cc}={n}" for cc, n in counts[:top])

    pages = 0
    print(f"[handles] start: {len(matched)} matched, {len(set(tld_map.values()))} countries, "
          f"cursor={after}", flush=True)
    while True:
        body = fetch_page(after)
        lines = [ln for ln in body.splitlines() if ln.strip()]
        if not lines:
            break
        last_created = after
        for ln in lines:
            try:
                rec = json.loads(ln)
            except json.JSONDecodeError:
                continue
            last_created = rec.get("createdAt") or last_created
            if rec.get("nullified"):
                continue
            did = rec.get("did")
            if not did:
                continue
            handle, tombstone = handle_of(rec)
            if tombstone:
                matched.pop(did, None)
                continue
            if handle is None:
                continue
            cc = country_of(handle, tld_map)
            if cc:
                matched[did] = cc      # add, or move between countries
            else:
                matched.pop(did, None)  # moved away from every mapped TLD
        pages += 1
        # No forward progress on the cursor -> we've caught up; stop.
        if last_created == after:
            break
        after = last_created
        if pages % args.flush_every == 0:
            by_country = flush()
            print(f"[handles] page {pages}, cursor {after}, {len(matched)} matched | "
                  f"{summary(by_country)}", flush=True)
        if len(lines) < 1000:
            break  # reached the end of the export
        if args.max_pages and pages >= args.max_pages:
            break
        if args.sleep:
            time.sleep(args.sleep)

    by_country = flush()
    print(f"[handles] done: {len(matched)} accounts across {len(by_country)} countries, "
          f"{pages} pages, cursor {after}", flush=True)
    print(f"[handles] {summary(by_country, top=99)}", flush=True)


if __name__ == "__main__":
    main()