"""Fetch author lists for every paper in the caption slice.

ArxivCap does not ship author names, so the author filter needs a one-time
metadata pass keyed on the arXiv IDs already in the slice. Uses the Semantic
Scholar batch endpoint (500 papers per request, ~110 requests for the astro
slice). Results append to authors.jsonl as they arrive, so an interrupted run
resumes, and the final parquet is rebuilt from that file.

    export S2_API_KEY=...        # optional but much faster
    python fetch_authors.py

Output: authors.parquet with columns (arxiv_id, authors, author_ids).
"""

import argparse
import json
import os
import time
from pathlib import Path

import duckdb
import pyarrow as pa
import pyarrow.parquet as pq
import requests

BATCH_URL = "https://api.semanticscholar.org/graph/v1/paper/batch"
BATCH_SIZE = 500


def load_ids(slice_path: str) -> list[str]:
    con = duckdb.connect()
    rows = con.execute(
        f"SELECT DISTINCT arxiv_id FROM read_parquet('{slice_path}') WHERE arxiv_id IS NOT NULL"
    ).fetchall()
    return [r[0] for r in rows]


def fetch_batch(ids: list[str], api_key: str) -> list:
    headers = {"x-api-key": api_key} if api_key else {}
    payload = {"ids": [f"ARXIV:{i}" for i in ids]}
    for attempt in range(6):
        try:
            r = requests.post(BATCH_URL, params={"fields": "authors"},
                              json=payload, headers=headers, timeout=60)
            if r.status_code == 429:
                wait = 10 * (attempt + 1)
                print(f"  rate limited, sleeping {wait}s")
                time.sleep(wait)
                continue
            r.raise_for_status()
            return r.json()
        except requests.exceptions.RequestException as e:
            if attempt == 5:
                raise
            time.sleep(5 * (attempt + 1))
    return []


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--slice", default="astro_captions.parquet")
    ap.add_argument("--out", default="authors.parquet")
    ap.add_argument("--cache", default="authors.jsonl")
    ap.add_argument("--sleep", type=float, default=1.0,
                    help="seconds between requests; raise if you see 429s")
    args = ap.parse_args()

    api_key = os.environ.get("S2_API_KEY", "")
    if not api_key:
        print("No S2_API_KEY set. This will work but is slower and rate limited.")

    cache_path = Path(args.cache)
    done = set()
    if cache_path.exists():
        for line in cache_path.open():
            done.add(json.loads(line)["arxiv_id"])

    ids = load_ids(args.slice)
    todo = [i for i in ids if i not in done]
    print(f"{len(ids):,} papers in slice, {len(done):,} already fetched, {len(todo):,} to go")

    with cache_path.open("a") as f:
        for start in range(0, len(todo), BATCH_SIZE):
            batch = todo[start:start + BATCH_SIZE]
            results = fetch_batch(batch, api_key)
            for arxiv_id, entry in zip(batch, results):
                authors, author_ids = [], []
                if entry and entry.get("authors"):
                    for a in entry["authors"]:
                        if a.get("name"):
                            authors.append(a["name"])
                            author_ids.append(a.get("authorId") or "")
                f.write(json.dumps({"arxiv_id": arxiv_id, "authors": authors,
                                    "author_ids": author_ids}) + "\n")
            f.flush()
            print(f"  {min(start + BATCH_SIZE, len(todo)):,}/{len(todo):,}")
            time.sleep(args.sleep)

    rows = [json.loads(l) for l in cache_path.open()]
    found = sum(1 for r in rows if r["authors"])
    pq.write_table(pa.table({
        "arxiv_id": [r["arxiv_id"] for r in rows],
        "authors": [r["authors"] for r in rows],
        "author_ids": [r["author_ids"] for r in rows],
    }), args.out)
    print(f"\n{len(rows):,} records written to {args.out}")
    print(f"{found:,} with at least one author ({100*found/max(len(rows),1):.1f}%)")


if __name__ == "__main__":
    main()