#!/usr/bin/env -S PYTHONUNBUFFERED=1 uv run --script --quiet
# /// script
# requires-python = ">=3.12"
# dependencies = []
# ///
"""
reconnect relay to all known PDS hosts.

fetches the community PDS list and sends requestCrawl for each host,
preventing the natural decay of connections as indie PDS hosts go quiet.

intended to run on a cron (every ~4 hours).

usage:
    ./scripts/reconnect
    ./scripts/reconnect --url https://relay.waow.tech --password "$RELAY_ADMIN_PASSWORD"
    ./scripts/reconnect --dry-run
"""

import argparse
import json
import sys
import time
import urllib.request


PDS_LIST_URL = "https://raw.githubusercontent.com/mary-ext/atproto-scraping/refs/heads/trunk/state.json"


def fetch_pds_list() -> list[str]:
    """fetch community-maintained PDS host list."""
    req = urllib.request.Request(PDS_LIST_URL)
    with urllib.request.urlopen(req, timeout=30) as resp:
        data = json.loads(resp.read())
    return [url.rstrip("/") for url in data.get("pdses", {}).keys() if url.startswith("https://")]


def request_crawl(relay_url: str, password: str, hostname: str) -> int:
    """send requestCrawl to relay admin endpoint. returns http status."""
    import base64
    payload = json.dumps({"hostname": hostname}).encode()
    auth = base64.b64encode(f"admin:{password}".encode()).decode()
    req = urllib.request.Request(
        f"{relay_url}/admin/pds/requestCrawl",
        data=payload,
        headers={"Content-Type": "application/json", "Authorization": f"Basic {auth}"},
        method="POST",
    )
    try:
        with urllib.request.urlopen(req, timeout=10) as resp:
            return resp.status
    except urllib.error.HTTPError as e:
        return e.code
    except (ConnectionError, OSError, urllib.error.URLError):
        return 0


def main():
    parser = argparse.ArgumentParser(description="reconnect relay to all known PDS hosts")
    parser.add_argument("--url", default="https://relay.waow.tech", help="relay URL")
    parser.add_argument("--password", default=None, help="relay admin password (or RELAY_ADMIN_PASSWORD env)")
    parser.add_argument("--delay", type=float, default=0.1, help="delay between requests in seconds")
    parser.add_argument("--dry-run", action="store_true", help="just fetch and count, don't send requests")
    args = parser.parse_args()

    import os
    password = args.password or os.environ.get("RELAY_ADMIN_PASSWORD", "")
    if not password and not args.dry_run:
        print("error: --password or RELAY_ADMIN_PASSWORD required", file=sys.stderr)
        return

    print(f"fetching PDS list from {PDS_LIST_URL}...")
    hosts = fetch_pds_list()
    print(f"found {len(hosts)} PDS hosts")

    if args.dry_run:
        print("dry run, exiting")
        return

    ok = 0
    errors = 0
    start = time.time()

    for i, host in enumerate(hosts):
        status = request_crawl(args.url, password, host)
        if status == 200:
            ok += 1
        else:
            errors += 1
            if errors <= 10:
                print(f"  error: {host} -> {status}")
            elif errors == 11:
                print(f"  (suppressing further errors)")

        if (i + 1) % 500 == 0:
            elapsed = time.time() - start
            print(f"  {i + 1}/{len(hosts)} ({ok} ok, {errors} errors, {elapsed:.0f}s)")

        time.sleep(args.delay)

    elapsed = time.time() - start
    print(f"done: {ok} ok, {errors} errors, {elapsed:.0f}s")


if __name__ == "__main__":
    main()