#!/usr/bin/env -S uv run --script --quiet
# /// script
# requires-python = ">=3.12"
# dependencies = [
#     "httpx",
#     "pydantic-settings",
#     "atproto",
# ]
# ///
"""
Backfill publication records via CAR walk.

Sibling of backfill-subscriptions. Our `publications` table is populated ONLY
when we directly witness a publication record (its create/update on the
firehose, or a backfill). Indexing a publication's *documents* does NOT
materialize the publication row. So any publication record that predates our
consumption window and hasn't been edited since has its documents indexed but
no publication row — and a subscription pointing at it finds no JOIN match.

That's a ~25% gap (measured: 393 / 1593 distinct subscribed-to publications had
no row), which silently understates the publisher and most-subscribed
leaderboards. This walks the two publication collections directly and upserts
the records ground-truth.

  com.atproto.sync.listReposByCollection -> DIDs holding the collection
  slingshot resolveMiniDoc                -> DID -> PDS endpoint
  com.atproto.sync.getRepo                -> full repo as CAR
  walk MST locally                        -> each publication record
  UPSERT into publications                -> idempotent on (uri)

indexed_at is set to now() on every upsert so the frozen local read-replica's
incremental sync (indexed_at >= last_sync) picks the rows up and rebuilds its
publications_fts — same contract every Turso-mutating script honors.

Usage:
    TURSO_URL=libsql://... TURSO_TOKEN=... ./scripts/backfill-publications
    ./scripts/backfill-publications --dry-run
    ./scripts/backfill-publications --workers 16
    ./scripts/backfill-publications --limit 50         # debug
"""

from __future__ import annotations

import argparse
import os
import sys
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass
from typing import NamedTuple

import httpx
from atproto_core.car import CAR
from atproto_core.cid import CID
from pydantic_settings import BaseSettings, SettingsConfigDict

SLINGSHOT = "https://slingshot.microcosm.blue"

# corpus admission policy — mirror of backend/src/policy.zig + ingest/ingester.zig
# isBridgyPds. Every write path (ingester, indexer, snapshot builder) drops
# these; a backfill must enforce the same gates or it re-admits exactly what we
# deliberately keep out. Registry of who/why/evidence: docs/exclusions.md.
BANNED_DIDS = frozenset(
    {
        "did:plc:oql6ds5vnff4ugar6rruliwd",  # drivepatents.com patent bot
        "did:plc:2s32mlusc66sjb256aenynfc",  # destinationcharged.com NHTSA recall mirror
    }
)


def is_bridgy_pds(pds: str | None) -> bool:
    """Host-match a PDS endpoint against brid.gy (mirrors backend isBridgyPds)."""
    if not pds:
        return False
    host = pds.split("://", 1)[-1].split("/", 1)[0].split(":", 1)[0].lower()
    return host == "brid.gy" or host.endswith(".brid.gy")


# both publication lexicons. leaflet stores the host in `base_path`; standard
# stores a full `url` we strip to its host (matches backend ingester.zig).
COLLECTIONS = ["pub.leaflet.publication", "site.standard.publication"]


class PubRow(NamedTuple):
    uri: str
    did: str
    rkey: str
    name: str
    description: str
    base_path: str

# Union DIDs across multiple relays — lightrail's index lags the Bluesky
# reference relay for niche collections.
RELAY_ENDPOINTS = [
    "https://relay1.us-east.bsky.network",
    "https://lightrail.microcosm.blue",
    "https://relay.waow.tech",
]


def strip_url_scheme(url: str | None) -> str | None:
    """site.standard.publication.url is a full URL (e.g. https://devlog.pckt.blog);
    we keep only the host, mirroring backend stripUrlScheme."""
    if not url:
        return None
    s = url
    for scheme in ("https://", "http://"):
        if s.startswith(scheme):
            s = s[len(scheme):]
            break
    host = s.split("/", 1)[0]
    return host or None


class Settings(BaseSettings):
    model_config = SettingsConfigDict(
        env_file=os.environ.get("ENV_FILE", ".env"), extra="ignore"
    )
    turso_url: str
    turso_token: str

    @property
    def turso_host(self) -> str:
        url = self.turso_url
        if url.startswith("libsql://"):
            url = url[len("libsql://"):]
        return url


# ---- Turso ----------------------------------------------------------------


def turso_pipeline(settings: Settings, requests: list[dict]) -> list[dict]:
    last_err: Exception | None = None
    for attempt in range(3):
        try:
            resp = httpx.post(
                f"https://{settings.turso_host}/v2/pipeline",
                headers={
                    "Authorization": f"Bearer {settings.turso_token}",
                    "Content-Type": "application/json",
                },
                json={"requests": [*requests, {"type": "close"}]},
                timeout=180,
            )
            resp.raise_for_status()
            return resp.json()["results"]
        except (httpx.ReadTimeout, httpx.ConnectError, httpx.RemoteProtocolError) as e:
            last_err = e
            print(
                f"  turso pipeline attempt {attempt + 1} failed: {type(e).__name__}; retrying",
                file=sys.stderr,
            )
            time.sleep(2 * (attempt + 1))
    assert last_err is not None
    raise last_err


def turso_upsert_publications(settings: Settings, rows: list[PubRow]) -> int:
    BATCH = 50
    written = 0
    for i in range(0, len(rows), BATCH):
        chunk = rows[i: i + BATCH]
        reqs = [
            {
                "type": "execute",
                "stmt": {
                    "sql": (
                        "INSERT INTO publications "
                        "(uri, did, rkey, name, description, base_path, indexed_at) "
                        "VALUES (?, ?, ?, ?, ?, ?, strftime('%Y-%m-%dT%H:%M:%S','now')) "
                        "ON CONFLICT(uri) DO UPDATE SET "
                        "  did = excluded.did, "
                        "  rkey = excluded.rkey, "
                        "  name = excluded.name, "
                        "  description = excluded.description, "
                        "  base_path = excluded.base_path, "
                        "  indexed_at = strftime('%Y-%m-%dT%H:%M:%S','now')"
                    ),
                    "args": [
                        {"type": "text", "value": row.uri},
                        {"type": "text", "value": row.did},
                        {"type": "text", "value": row.rkey},
                        {"type": "text", "value": row.name},
                        {"type": "text", "value": row.description},
                        {"type": "text", "value": row.base_path},
                    ],
                },
            }
            for row in chunk
        ]
        results = turso_pipeline(settings, reqs)
        for r in results:
            if r["type"] == "error":
                print(f"  upsert error: {r['error']}", file=sys.stderr)
                continue
            response = r.get("response") or {}
            if response.get("type") != "execute":
                continue
            written += int(response.get("result", {}).get("affected_row_count", 0))
    return written


# ---- Enumeration + resolution -------------------------------------------


def _enumerate_from_relay(client: httpx.Client, base: str, collection: str) -> list[str]:
    dids: list[str] = []
    cursor: str | None = None
    while True:
        params: dict = {"collection": collection, "limit": 1000}
        if cursor:
            params["cursor"] = cursor
        try:
            r = client.get(
                f"{base}/xrpc/com.atproto.sync.listReposByCollection",
                params=params,
                timeout=30,
            )
            r.raise_for_status()
        except Exception as e:
            print(f"  {base} ({collection}): relay error {type(e).__name__}", file=sys.stderr)
            return dids
        body = r.json()
        for repo in body.get("repos", []):
            dids.append(repo["did"])
        cursor = body.get("cursor")
        if not cursor:
            return dids


def enumerate_owner_dids() -> list[str]:
    all_dids: set[str] = set()
    with httpx.Client() as client:
        for collection in COLLECTIONS:
            for base in RELAY_ENDPOINTS:
                from_relay = _enumerate_from_relay(client, base, collection)
                print(f"  {base} ({collection}): {len(from_relay)} DIDs")
                all_dids.update(from_relay)
    return sorted(all_dids)


def _did_from_pub_uri(uri: str) -> str | None:
    # at://<did>/<collection>/<rkey>
    if not uri.startswith("at://"):
        return None
    rest = uri[len("at://"):]
    did = rest.split("/", 1)[0]
    return did or None


def owner_dids_with_unindexed_subscribed_pubs(settings: Settings) -> list[str]:
    """The surgical target set: owners of publications that someone subscribes to
    but which have no row in our publications table. Directly closes the measured
    JOIN gap with a few hundred CAR fetches instead of walking the whole network."""
    results = turso_pipeline(
        settings,
        [
            {
                "type": "execute",
                "stmt": {
                    "sql": (
                        "SELECT DISTINCT s.publication_uri "
                        "FROM subscriptions s "
                        "LEFT JOIN publications p ON p.uri = s.publication_uri "
                        "WHERE p.uri IS NULL"
                    )
                },
            }
        ],
    )
    dids: set[str] = set()
    for r in results:
        if r.get("type") != "ok":
            continue
        response = r.get("response") or {}
        if response.get("type") != "execute":
            continue
        for row in response.get("result", {}).get("rows", []):
            if not row:
                continue
            uri = row[0].get("value")
            did = _did_from_pub_uri(uri) if isinstance(uri, str) else None
            if did:
                dids.add(did)
    return sorted(dids)


def resolve_pds(client: httpx.Client, did: str) -> str | None:
    try:
        r = client.get(
            f"{SLINGSHOT}/xrpc/blue.microcosm.identity.resolveMiniDoc",
            params={"identifier": did},
            timeout=15,
        )
        r.raise_for_status()
        return r.json().get("pds")
    except Exception:
        return None


# ---- CAR walk ------------------------------------------------------------


def _walk_mst(car: CAR, prefixes: list[bytes]) -> list[tuple[str, str, dict]]:
    """Return [(collection, rkey, decoded_value_dict), ...] for every entry whose
    key starts with one of the f"{collection}/" prefixes."""
    blocks = car.blocks
    root = blocks[car.root]
    data_cid = CID.decode(root["data"])

    out: list[tuple[str, str, dict]] = []
    stack = [data_cid]
    while stack:
        node = blocks[stack.pop()]
        if (left := node.get("l")) is not None:
            stack.append(CID.decode(left))
        prev_key = b""
        for entry in node["e"]:
            key = prev_key[: entry["p"]] + entry["k"]
            prev_key = key
            for prefix in prefixes:
                if key.startswith(prefix):
                    value_cid = CID.decode(entry["v"])
                    value = blocks[value_cid]
                    collection = prefix[:-1].decode()  # strip trailing "/"
                    rkey = key[len(prefix):].decode("utf-8", errors="replace")
                    out.append((collection, rkey, value))
                    break
            if (t := entry.get("t")) is not None:
                stack.append(CID.decode(t))
    return out


# ---- Fetch one DID's publications ----------------------------------------


@dataclass
class FetchResult:
    did: str
    rows: list[PubRow]
    status: str  # "ok", "no-pds", "fetch-err:<type>", "parse-err"


def fetch_publications_for_did(client: httpx.Client, did: str) -> FetchResult:
    # policy gate 1: explicitly banned bulk-archive repos (docs/exclusions.md).
    if did in BANNED_DIDS:
        return FetchResult(did, [], "banned")
    pds = resolve_pds(client, did)
    if not pds:
        return FetchResult(did, [], "no-pds")
    # policy gate 2: bridgy-hosted PDS — same drop as every other write path.
    if is_bridgy_pds(pds):
        return FetchResult(did, [], "bridgy")
    try:
        r = client.get(
            f"{pds}/xrpc/com.atproto.sync.getRepo",
            params={"did": did},
            timeout=60,
        )
        r.raise_for_status()
        car = CAR.from_bytes(r.content)
    except httpx.HTTPStatusError as e:
        return FetchResult(did, [], f"fetch-err:{e.response.status_code}")
    except Exception as e:
        return FetchResult(did, [], f"fetch-err:{type(e).__name__}")
    try:
        prefixes = [f"{c}/".encode() for c in COLLECTIONS]
        entries = _walk_mst(car, prefixes)
    except Exception:
        return FetchResult(did, [], "parse-err")

    rows: list[PubRow] = []
    for collection, rkey, value in entries:
        if not isinstance(value, dict):
            continue
        name = value.get("name")
        if not isinstance(name, str) or not name:
            continue  # name is required (matches backend processPublication)
        description = value.get("description")
        description = description if isinstance(description, str) else ""
        # leaflet: base_path; standard: host of url
        base_path = value.get("base_path")
        if not isinstance(base_path, str) or not base_path:
            base_path = strip_url_scheme(value.get("url")) or ""
        # skip dev/staging .test domains (matches backend)
        if base_path.endswith(".test"):
            continue
        rows.append(
            PubRow(
                uri=f"at://{did}/{collection}/{rkey}",
                did=did,
                rkey=rkey,
                name=name,
                description=description,
                base_path=base_path,
            )
        )
    return FetchResult(did, rows, "ok")


# ---- Main ---------------------------------------------------------------


def main() -> int:
    parser = argparse.ArgumentParser(description="Backfill publication records via CAR walk")
    parser.add_argument("--dry-run", action="store_true", help="Don't write to Turso")
    parser.add_argument("--workers", type=int, default=8, help="Concurrent PDS fetches")
    parser.add_argument("--limit", type=int, default=None, help="Cap DIDs (debug)")
    parser.add_argument(
        "--from-subscriptions",
        action="store_true",
        help="Surgical mode: only walk owners of subscribed-but-unindexed pubs "
        "(closes the measured JOIN gap with a few hundred fetches instead of the "
        "whole network).",
    )
    args = parser.parse_args()

    try:
        settings = Settings()  # type: ignore
    except Exception as e:
        print(f"error loading settings: {e}", file=sys.stderr)
        print("required env vars: TURSO_URL, TURSO_TOKEN", file=sys.stderr)
        return 1

    if args.from_subscriptions:
        print("=== surgical: owners of subscribed-but-unindexed publications ===")
        dids = owner_dids_with_unindexed_subscribed_pubs(settings)
    else:
        print(f"=== {' + '.join(COLLECTIONS)} ===")
        print("enumerating owner DIDs...")
        dids = enumerate_owner_dids()
    if args.limit:
        dids = dids[: args.limit]
    print(f"  {len(dids)} DIDs")
    if not dids:
        return 0

    started = time.time()
    done = 0
    statuses: dict[str, int] = {}
    rows_acc: list[PubRow] = []

    # follow_redirects: some PDSes 301 getRepo to a canonical host.
    with httpx.Client(follow_redirects=True) as client, ThreadPoolExecutor(max_workers=args.workers) as pool:
        futures = {pool.submit(fetch_publications_for_did, client, did): did for did in dids}
        for fut in as_completed(futures):
            res = fut.result()
            statuses[res.status] = statuses.get(res.status, 0) + 1
            rows_acc.extend(res.rows)
            done += 1
            if done % 25 == 0 or done == len(dids):
                elapsed = time.time() - started
                rate = done / elapsed if elapsed > 0 else 0
                print(f"  [{done}/{len(dids)}] {rate:.1f} dids/s, {len(rows_acc)} records, statuses={statuses}")

    print(f"collected {len(rows_acc)} records from {statuses.get('ok', 0)} DIDs")
    if args.dry_run:
        print("dry-run: skipping writes")
        return 0
    print("upserting...")
    written = turso_upsert_publications(settings, rows_acc)
    print(f"  wrote {written} rows")
    return 0


if __name__ == "__main__":
    sys.exit(main())