#!/usr/bin/env python3
"""
s2_citation_crawler.py

BFS citation graph crawler for Semantic Scholar via the s2-proxy Cloudflare
Worker fleet.  Feeds the result into pagerank_eigenvalue_survey.run_survey().

Architecture:
  coordinator (this script)
    └─ round-robins requests across N worker URLs
         each worker: KV cache → S2 API (with per-worker key) → KV cache write

Usage:
  # Single worker (no key, slow but free):
  python s2_citation_crawler.py --seed "PageRank Brin Page 1998" --hops 2

  # Worker fleet (fast):
  python s2_citation_crawler.py \\
    --seed "PageRank Brin Page 1998" \\
    --hops 2 \\
    --workers https://s2-proxy.YOUR.workers.dev https://s2-proxy-2.YOUR.workers.dev \\
    --out citation_graph.json

Requires: requests, networkx, scipy, numpy
  uv pip install requests networkx scipy numpy
"""

import argparse
import json
import time
import sys
from collections import deque
from dataclasses import dataclass, asdict, field
from pathlib import Path
from typing import List, Dict, Optional, Set, Iterator
from itertools import cycle

import requests
import networkx as nx

# ── Import spectral survey from sibling script ────────────────────────────────
sys.path.insert(0, str(Path(__file__).parent))
try:
    from pagerank_eigenvalue_survey import run_survey, BIN_COUNT, verify_spectral_gap
    HAS_SURVEY = True
except ImportError:
    HAS_SURVEY = False
    print("[warn] pagerank_eigenvalue_survey not found — graph will be saved but not analysed")

# ── S2 field set ──────────────────────────────────────────────────────────────

PAPER_FIELDS = "paperId,title,year,referenceCount,citationCount,influentialCitationCount"
REF_FIELDS   = "paperId,title,year,citationCount"
CITE_FIELDS  = "paperId,title,year,citationCount"

# ── Data types ────────────────────────────────────────────────────────────────

@dataclass
class S2Paper:
    paper_id: str
    title: str
    year: Optional[int]
    ref_count: int
    cite_count: int
    influential_cite_count: int
    references: List[str] = field(default_factory=list)   # paper_ids
    citations:  List[str] = field(default_factory=list)   # paper_ids

# ── Worker fleet client ───────────────────────────────────────────────────────

class WorkerFleet:
    """Round-robins requests across the s2-proxy worker fleet."""

    DIRECT_BASE = "https://api.semanticscholar.org"

    def __init__(self, worker_urls: List[str], api_key: Optional[str] = None,
                 rate_delay: float = 1.1):
        self.worker_urls  = worker_urls
        self.api_key      = api_key
        self.rate_delay   = rate_delay  # seconds between requests (per worker slot)
        self._fleet       = cycle(worker_urls) if worker_urls else None
        self._last_req    = 0.0
        self._stats       = {"hits": 0, "misses": 0, "errors": 0, "requests": 0}

    def _next_url_base(self) -> str:
        if self._fleet:
            return next(self._fleet)
        return self.DIRECT_BASE

    def _throttle(self):
        gap = time.monotonic() - self._last_req
        if gap < self.rate_delay:
            time.sleep(self.rate_delay - gap)
        self._last_req = time.monotonic()

    def get(self, s2_path: str, params: Dict = None) -> Optional[dict]:
        self._throttle()
        base = self._next_url_base()
        self._stats["requests"] += 1

        try:
            if self._fleet:
                # Via worker
                worker_params = {"path": s2_path, **(params or {})}
                resp = requests.get(f"{base}/fetch", params=worker_params, timeout=15)
                cache_status = resp.headers.get("X-S2-Cache", "?")
                if cache_status == "HIT":
                    self._stats["hits"] += 1
                else:
                    self._stats["misses"] += 1
            else:
                # Direct S2
                headers = {}
                if self.api_key:
                    headers["x-api-key"] = self.api_key
                resp = requests.get(
                    f"{self.DIRECT_BASE}{s2_path}",
                    params=params or {},
                    headers=headers,
                    timeout=15,
                )

            # Back-off on rate limit
            if resp.status_code == 429:
                retry_after = int(resp.headers.get("retry-after", 60))
                print(f"  [rate-limit] waiting {retry_after}s")
                time.sleep(retry_after)
                return self.get(s2_path, params)  # retry once

            resp.raise_for_status()
            return resp.json()

        except Exception as exc:
            self._stats["errors"] += 1
            print(f"  [fetch error] {s2_path}: {exc}")
            return None

    def stats(self) -> str:
        s = self._stats
        hit_rate = (s["hits"] / max(s["requests"], 1)) * 100
        return (f"requests={s['requests']}  cache_hits={s['hits']}({hit_rate:.0f}%)"
                f"  misses={s['misses']}  errors={s['errors']}")


# ── S2 API wrappers ───────────────────────────────────────────────────────────

def search_paper(fleet: WorkerFleet, query: str) -> Optional[str]:
    """Return the top S2 paper_id for a title/author query."""
    data = fleet.get("/graph/v1/paper/search", {
        "query": query, "fields": "paperId,title,year", "limit": 1
    })
    if not data or not data.get("data"):
        return None
    hit = data["data"][0]
    print(f"  seed paper: [{hit.get('year','')}] {hit.get('title','?')}")
    return hit["paperId"]


def fetch_paper(fleet: WorkerFleet, paper_id: str) -> Optional[S2Paper]:
    data = fleet.get(f"/graph/v1/paper/{paper_id}", {"fields": PAPER_FIELDS})
    if not data or "paperId" not in data:
        return None
    return S2Paper(
        paper_id=data["paperId"],
        title=data.get("title") or "",
        year=data.get("year"),
        ref_count=data.get("referenceCount", 0),
        cite_count=data.get("citationCount", 0),
        influential_cite_count=data.get("influentialCitationCount", 0),
    )


def fetch_neighbors(fleet: WorkerFleet, paper_id: str,
                    max_refs: int = 50, max_cites: int = 50
                    ) -> tuple[List[str], List[str]]:
    """Return (reference_ids, citation_ids) for a paper, capped at max_*."""
    refs, cites = [], []

    ref_data = fleet.get(f"/graph/v1/paper/{paper_id}/references", {
        "fields": REF_FIELDS, "limit": max_refs
    })
    if ref_data and "data" in ref_data:
        refs = [e["citedPaper"]["paperId"]
                for e in ref_data["data"]
                if e.get("citedPaper", {}).get("paperId")]

    cite_data = fleet.get(f"/graph/v1/paper/{paper_id}/citations", {
        "fields": CITE_FIELDS, "limit": max_cites
    })
    if cite_data and "data" in cite_data:
        cites = [e["citingPaper"]["paperId"]
                 for e in cite_data["data"]
                 if e.get("citingPaper", {}).get("paperId")]

    return refs, cites


# ── BFS crawler ───────────────────────────────────────────────────────────────

def crawl(fleet: WorkerFleet, seed_id: str, hops: int,
          max_refs: int = 40, max_cites: int = 40,
          max_papers: int = 500) -> Dict[str, S2Paper]:
    """BFS from seed_id for `hops` hops. Returns dict of paper_id → S2Paper."""
    visited: Set[str] = set()
    papers:  Dict[str, S2Paper] = {}
    queue:   deque = deque([(seed_id, 0)])

    print(f"\nBFS crawl: seed={seed_id}  hops={hops}  max_papers={max_papers}")

    while queue and len(papers) < max_papers:
        pid, depth = queue.popleft()
        if pid in visited:
            continue
        visited.add(pid)

        paper = fetch_paper(fleet, pid)
        if not paper:
            continue
        papers[pid] = paper
        print(f"  [{depth}/{hops}] {len(papers):>4} papers  "
              f"cite={paper.cite_count}  {paper.title[:60]}")

        if depth < hops:
            refs, cites = fetch_neighbors(fleet, pid, max_refs, max_cites)
            paper.references = refs
            paper.citations  = cites
            for nid in refs + cites:
                if nid not in visited:
                    queue.append((nid, depth + 1))

    print(f"\nCrawl complete: {len(papers)} papers, {fleet.stats()}")
    return papers


# ── Build NetworkX graph ──────────────────────────────────────────────────────

def build_graph(papers: Dict[str, S2Paper]) -> nx.DiGraph:
    """Citation graph: edge A→B means A cites B."""
    g = nx.DiGraph()
    pid_to_int = {pid: i for i, pid in enumerate(papers)}

    for pid, paper in papers.items():
        node_id = pid_to_int[pid]
        g.add_node(node_id,
                   paper_id=pid,
                   title=paper.title[:80],
                   year=paper.year or 0,
                   cite_count=paper.cite_count)

    for pid, paper in papers.items():
        src = pid_to_int[pid]
        for ref_id in paper.references:
            if ref_id in pid_to_int:
                g.add_edge(src, pid_to_int[ref_id])  # src cites ref

    print(f"Graph: {g.number_of_nodes()} nodes, {g.number_of_edges()} edges")
    return g


# ── Main ──────────────────────────────────────────────────────────────────────

def main():
    ap = argparse.ArgumentParser(description="S2 citation BFS crawler + spectral survey")
    ap.add_argument("--seed", default="PageRank Brin Page 1998 web search engine",
                    help="Title/author query to find seed paper")
    ap.add_argument("--seed-id", default=None,
                    help="Use a specific S2 paper_id instead of title search")
    ap.add_argument("--hops", type=int, default=2,
                    help="BFS depth (default 2; 3 gets very large)")
    ap.add_argument("--max-refs", type=int, default=40,
                    help="Max references to follow per paper")
    ap.add_argument("--max-cites", type=int, default=40,
                    help="Max citations to follow per paper")
    ap.add_argument("--max-papers", type=int, default=300,
                    help="Hard cap on total papers fetched")
    ap.add_argument("--workers", nargs="*", default=[],
                    help="s2-proxy worker base URLs (omit to hit S2 directly)")
    ap.add_argument("--api-key", default=None,
                    help="S2 API key for direct mode (not used in worker mode)")
    ap.add_argument("--rate-delay", type=float, default=1.1,
                    help="Seconds between requests per worker slot")
    ap.add_argument("--out", default="citation_graph.json",
                    help="Output file for the paper graph JSON")
    ap.add_argument("--no-survey", action="store_true",
                    help="Skip spectral survey (just crawl and save)")
    args = ap.parse_args()

    fleet = WorkerFleet(args.workers, api_key=args.api_key, rate_delay=args.rate_delay)

    # Resolve seed
    if args.seed_id:
        seed_id = args.seed_id
    else:
        print(f"Searching for seed: {args.seed}")
        seed_id = search_paper(fleet, args.seed)
        if not seed_id:
            print("ERROR: seed paper not found")
            sys.exit(1)
    print(f"Seed paper_id: {seed_id}")

    # Crawl
    papers = crawl(fleet, seed_id,
                   hops=args.hops,
                   max_refs=args.max_refs,
                   max_cites=args.max_cites,
                   max_papers=args.max_papers)

    # Save
    out = Path(args.out)
    with out.open("w") as f:
        json.dump({pid: asdict(p) for pid, p in papers.items()}, f, indent=2)
    print(f"Saved {len(papers)} papers → {out}")

    # Build graph and run spectral survey
    g = build_graph(papers)
    if HAS_SURVEY and not args.no_survey:
        print("\n" + "=" * 72)
        print("Running spectral survey on real S2 citation graph...")
        print("=" * 72)
        run_survey(g)
    else:
        # Basic stats without survey
        pr = nx.pagerank(g, alpha=0.85)
        top10 = sorted(pr.items(), key=lambda x: -x[1])[:10]
        print("\nTop 10 by PageRank:")
        for nid, score in top10:
            title = g.nodes[nid].get("title", "?")
            year  = g.nodes[nid].get("year", "?")
            print(f"  [{year}] {score:.5f}  {title}")


if __name__ == "__main__":
    main()