#!/usr/bin/env python3
"""Tiny Semantic Scholar Graph API client for the literature search (stdlib only).

Every API response is cached under ../search/.cache/ so re-running a search is
free and reproducible, and the raw JSON is kept as an audit trail.

Usage:
    python3 s2.py search "query string" [--year 2020-] [--limit 100] [--out NAME]
    python3 s2.py refs   ARXIV_ID_OR_S2_ID [--out NAME]   # backward snowball
    python3 s2.py cites  ARXIV_ID_OR_S2_ID [--out NAME]   # forward snowball

Output: a CSV in ../search/results/ (one row per paper) and a line appended
to ../search/search-log.csv recording what was run, when, and how many hits.

Set S2_API_KEY in the environment if you get one (free, raises rate limits).
"""
import csv
import datetime as dt
import hashlib
import json
import os
import sys
import time
import urllib.error
import urllib.parse
import urllib.request

HERE = os.path.dirname(os.path.abspath(__file__))
SEARCH_DIR = os.path.normpath(os.path.join(HERE, "..", "search"))
CACHE = os.path.join(SEARCH_DIR, ".cache")
RESULTS = os.path.join(SEARCH_DIR, "results")
LOG = os.path.join(SEARCH_DIR, "search-log.csv")
API = "https://api.semanticscholar.org/graph/v1"
FIELDS = "paperId,externalIds,title,year,venue,authors,citationCount,abstract,publicationDate"


def get(path, params):
    url = f"{API}/{path}?{urllib.parse.urlencode(params)}"
    key = hashlib.sha1(url.encode()).hexdigest()
    cpath = os.path.join(CACHE, key + ".json")
    if os.path.exists(cpath):
        return json.load(open(cpath))
    headers = {"User-Agent": "msc-thesis-lit-search/0.1"}
    if os.environ.get("S2_API_KEY"):
        headers["x-api-key"] = os.environ["S2_API_KEY"]
    for attempt in range(int(os.environ.get("S2_RETRIES", 8))):
        try:
            with urllib.request.urlopen(urllib.request.Request(url, headers=headers), timeout=60) as r:
                data = json.load(r)
            os.makedirs(CACHE, exist_ok=True)
            json.dump(data, open(cpath, "w"))
            time.sleep(1.2)  # stay under the shared rate limit
            return data
        except urllib.error.HTTPError as e:
            if e.code == 404:
                return None
            if e.code not in (429, 500, 502, 503, 504):
                raise
        except (urllib.error.URLError, TimeoutError):
            pass
        wait = min(60, 3 * 2 ** attempt)
        print(f"  rate-limited/transient error; retry in {wait}s", file=sys.stderr)
        time.sleep(wait)
    sys.exit(f"gave up on {url}")


def norm_id(pid):
    if pid[:1].isdigit() and "." in pid:  # bare arXiv id like 2309.14322
        return "arXiv:" + pid
    return pid


def row(p):
    ext = p.get("externalIds") or {}
    return {
        "s2_id": p.get("paperId"),
        "arxiv": ext.get("ArXiv", ""),
        "doi": ext.get("DOI", ""),
        "year": p.get("year") or "",
        "title": " ".join((p.get("title") or "").split()),
        "first_author": ((p.get("authors") or [{}])[0] or {}).get("name", ""),
        "n_authors": len(p.get("authors") or []),
        "venue": p.get("venue") or "",
        "citations": p.get("citationCount") or 0,
        "abstract": " ".join((p.get("abstract") or "").split()),
    }


def search(query, year=None, limit=100):
    """Relevance-ranked search (top `limit` hits)."""
    out, offset, total = [], 0, 0
    while len(out) < limit:
        params = {"query": query, "fields": FIELDS, "offset": offset, "limit": min(100, limit - len(out))}
        if year:
            params["year"] = year
        data = get("paper/search", params)
        out += data.get("data") or []
        total = data.get("total", len(out))
        if "next" not in data:
            break
        offset = data["next"]
    return out[:limit], total


def snowball(pid, direction):
    field = "citedPaper" if direction == "refs" else "citingPaper"
    path = f"paper/{norm_id(pid)}/{'references' if direction == 'refs' else 'citations'}"
    out, offset = [], 0
    while True:
        data = get(path, {"fields": FIELDS, "limit": 500, "offset": offset})
        if data is None:
            sys.exit(f"paper not found: {pid}")
        out += [d[field] for d in data.get("data") or [] if d.get(field, {}).get("title")]
        if "next" not in data:
            return out, len(out)
        offset = data["next"]


def write(name, papers, what, total):
    os.makedirs(RESULTS, exist_ok=True)
    rows = [row(p) for p in papers]
    path = os.path.join(RESULTS, name + ".csv")
    with open(path, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=list(rows[0]) if rows else ["s2_id"])
        w.writeheader()
        w.writerows(rows)
    new = not os.path.exists(LOG)
    with open(LOG, "a", newline="") as f:
        w = csv.writer(f)
        if new:
            w.writerow(["date", "source", "action", "query_or_seed", "filters", "total_hits", "retrieved", "results_file"])
        w.writerow([dt.date.today().isoformat(), "Semantic Scholar API", *what, total, len(rows), os.path.relpath(path, SEARCH_DIR)])
    print(f"{what[0]} '{what[1]}': {total} total, {len(rows)} saved -> {os.path.relpath(path, SEARCH_DIR)}")


def main(a):
    def opt(flag, default=None):
        return a[a.index(flag) + 1] if flag in a else default

    cmd, arg = a[0], a[1]
    out = opt("--out", hashlib.sha1(f"{cmd}{arg}".encode()).hexdigest()[:8])
    if cmd == "search":
        year = opt("--year")
        papers, total = search(arg, year, int(opt("--limit", 100)))
        write(out, papers, ("search", arg, f"year={year or 'any'}"), total)
    elif cmd in ("refs", "cites"):
        papers, total = snowball(arg, cmd)
        label = "backward snowball (references)" if cmd == "refs" else "forward snowball (citations)"
        write(out, papers, (label, arg, ""), total)
    else:
        sys.exit(__doc__)


if __name__ == "__main__":
    main(sys.argv[1:])