src/step_01_fetch/fetch_metadata.py

from pathlib import Path
import argparse
import json
import os
import random
import re
import tarfile
import time

import arxiv
import requests

from config import ACL_IDS_PATH
from process_tex_source import preprocess_tex, extract_introduction_and_related
from semanticscholar_client import get_paper, get_paper_links, search_by_title


def load_ids(path: Path):
    return json.loads(path.read_text(encoding="utf-8"))


def ensure_dir(path: Path):
    path.mkdir(parents=True, exist_ok=True)


_ARXIV_LAST_TS = 0.0


def _cleanup_partial_source_dir(source_dir: Path) -> None:
    for pattern in ("*.tar.gz", "*.tgz", "*.tar"):
        for path in source_dir.glob(pattern):
            try:
                path.unlink()
            except Exception:
                pass


def _download_arxiv_source_with_retries(paper, source_dir: Path, arxiv_id: str) -> Path | None:
    max_retries = int(os.getenv("ARXIV_SOURCE_MAX_RETRIES", "4"))
    base_sleep = float(os.getenv("ARXIV_SOURCE_BASE_SLEEP", "2.0"))
    max_sleep = float(os.getenv("ARXIV_MAX_BACKOFF", "60"))
    last_exc = None

    for attempt in range(max_retries):
        _cleanup_partial_source_dir(source_dir)
        try:
            _arxiv_min_interval_sleep()
            tar_path = Path(paper.download_source(dirpath=str(source_dir)))
            if not tar_path.exists():
                raise FileNotFoundError(f"download_source returned {tar_path}, but the file does not exist")
            if tar_path.stat().st_size < 1024:
                raise IOError(f"downloaded source archive is unexpectedly small ({tar_path.stat().st_size} bytes)")
            return tar_path
        except Exception as exc:
            last_exc = exc
            sleep = min(base_sleep * (2**attempt), max_sleep) + random.uniform(0.0, 0.5)
            print(f"[WARN] Failed to download source for {arxiv_id} on attempt {attempt + 1}/{max_retries}: {exc}")
            if attempt + 1 < max_retries:
                print(f"[INFO] Retrying source download in {sleep:.2f}s")
                time.sleep(sleep)

    print(f"[WARN] Source download failed for {arxiv_id} after {max_retries} attempts: {last_exc}")
    return None


def _arxiv_min_interval_sleep() -> None:
    """Global throttle to avoid arXiv API rate limits."""
    global _ARXIV_LAST_TS
    min_interval = float(os.getenv("ARXIV_MIN_INTERVAL", "1.0"))
    now = time.monotonic()
    elapsed = now - _ARXIV_LAST_TS
    if elapsed < min_interval:
        time.sleep(min_interval - elapsed)
    _ARXIV_LAST_TS = time.monotonic()


def download_arxiv_tex(arxiv_id: str, base_dir: Path) -> Path | None:
    """
    Download LaTeX source from arXiv and return the path to a merged TeX file.

    - arxiv_id: e.g. "2410.22815"
    - base_dir: paper directory where source should be unpacked
    """
    source_dir = base_dir / f"tex_{arxiv_id}"
    source_dir.mkdir(parents=True, exist_ok=True)
    search = arxiv.Search(id_list=[arxiv_id])
    max_retries = int(os.getenv("ARXIV_MAX_RETRIES", "6"))
    base_sleep = float(os.getenv("ARXIV_BASE_SLEEP", "2.0"))
    max_sleep = float(os.getenv("ARXIV_MAX_BACKOFF", "60"))
    paper = None

    for attempt in range(max_retries):
        try:
            _arxiv_min_interval_sleep()
            paper = next(search.results())
            break
        except StopIteration:
            print(f"[WARN] No arXiv paper found for ID {arxiv_id}")
            return None
        except arxiv.HTTPError as exc:
            if getattr(exc, "status", None) == 429 or "429" in str(exc):
                sleep = min(base_sleep * (2**attempt), max_sleep) + random.uniform(0.0, 0.5)
                print(f"[WARN] arXiv 429 → retrying in {sleep:.2f}s")
                time.sleep(sleep)
                continue
            print(f"[WARN] arXiv HTTP error for {arxiv_id}: {exc}")
            return None
        except Exception as exc:
            sleep = min(base_sleep * (2**attempt), max_sleep) + random.uniform(0.0, 0.5)
            print(f"[WARN] arXiv error {exc} → retrying in {sleep:.2f}s")
            time.sleep(sleep)
            continue

    if paper is None:
        print(f"[ERROR] Giving up after {max_retries} attempts for arXiv ID {arxiv_id}")
        return None

    tar_path = _download_arxiv_source_with_retries(paper, source_dir, arxiv_id)
    if tar_path is None:
        return None

    try:
        with tarfile.open(tar_path) as tar:
            tar.extractall(path=source_dir)
        os.remove(tar_path)
    except Exception as exc:
        print(f"[WARN] Failed to extract source for {arxiv_id}: {exc}")
        return None

    processed_tex = preprocess_tex(source_dir)
    if processed_tex:
        extract_introduction_and_related(processed_tex)

    if not processed_tex or not processed_tex.exists():
        print(f"[WARN] Could not produce merged TeX for {arxiv_id}")
        return None

    print(f"[INFO] Processed LaTeX for {arxiv_id} at {processed_tex}")
    return processed_tex


def _extract_arxiv_id_from_text(text: str) -> str | None:
    if not text:
        return None
    match = re.search(r"\b(\d{4}\.\d{4,5}(?:v\d+)?)\b", text)
    if match:
        return match.group(1)
    match = re.search(r"arxiv[:\s/]*(\d{4}\.\d{4,5}(?:v\d+)?)", text, re.IGNORECASE)
    if match:
        return match.group(1)
    return None


def _safe_write_json(path: Path, payload) -> None:
    path.write_text(json.dumps(payload, indent=2), encoding="utf-8")


def _safe_write_text(path: Path, text: str) -> None:
    path.write_text(text, encoding="utf-8")


def _query_openreview_for_paper(openreview_id: str) -> dict | None:
    """Query OpenReview using a real OpenReview note/forum id."""
    if not openreview_id:
        return None

    try_urls = [
        f"https://api.openreview.net/notes?forum={openreview_id}",
        f"https://api2.openreview.net/notes?forum={openreview_id}",
        f"https://api.openreview.net/notes?id={openreview_id}",
        f"https://api2.openreview.net/notes?id={openreview_id}",
    ]

    for url in try_urls:
        try:
            response = requests.get(url, timeout=20)
            if response.status_code != 200:
                continue
            payload = response.json()
        except Exception:
            continue

        notes = None
        if isinstance(payload, dict) and isinstance(payload.get("notes"), list):
            notes = payload["notes"]
        elif isinstance(payload, dict) and payload.get("content"):
            notes = [payload]
        elif isinstance(payload, list):
            notes = payload

        if not notes:
            continue

        note = notes[0]
        content = note.get("content") if isinstance(note, dict) else None
        title = None
        arxiv_id = None
        pdf_url = None

        if isinstance(content, dict):
            raw_title = content.get("title") or content.get("paperTitle")
            title = raw_title.get("value") if isinstance(raw_title, dict) else raw_title

            raw_pdf = content.get("pdf")
            pdf_url = raw_pdf.get("value") if isinstance(raw_pdf, dict) else raw_pdf

            for value in content.values():
                if isinstance(value, dict):
                    value = value.get("value")
                if isinstance(value, list):
                    value = " ".join(str(item) for item in value)
                if isinstance(value, str):
                    arxiv_id = _extract_arxiv_id_from_text(value)
                    if arxiv_id:
                        break

        if not title and isinstance(note, dict):
            title = note.get("title") or note.get("forumTitle")

        if not arxiv_id and isinstance(note, dict):
            for value in note.values():
                if isinstance(value, str):
                    arxiv_id = _extract_arxiv_id_from_text(value)
                    if arxiv_id:
                        break

        return {
            "title": title,
            "arxiv_id": arxiv_id,
            "pdf_url": pdf_url,
            "openreview_id": openreview_id,
            "source_url": url,
        }

    return None


def _treat_as_openreview(paper: dict) -> bool:
    acl_id = str(paper.get("id", "")).lower()
    id_type = str(paper.get("id_type", "")).lower()
    return (
        id_type == "openreview"
        or bool(paper.get("openreview_id"))
        or acl_id.startswith("neurips-")
        or acl_id.startswith("icml-")
    )


def _fetch_s2_by_title(title: str, acl_id: str) -> tuple[int, dict | None]:
    if not title:
        print(f"[WARN] no title available for {acl_id} → skipping.")
        return 0, None
    hit = search_by_title(title)
    if not hit:
        print(f"[WARN] no S2 match for {acl_id} ({title}) → skipping.")
        return 0, None
    s2_id = hit["paperId"]
    print(f"[DEBUG] title search matched semantic scholar paperId={s2_id}")
    return get_paper(s2_id, id_type="SemanticScholar")


def _best_arxiv_id(*values: str) -> str | None:
    for value in values:
        arxiv_id = _extract_arxiv_id_from_text(value or "")
        if arxiv_id:
            return arxiv_id
    return None


def _write_openreview_snapshot(paper_dir: Path, payload: dict) -> None:
    if payload:
        _safe_write_json(paper_dir / "openreview_metadata.json", payload)


def _write_metadata_outputs(paper_dir: Path, acl_id: str, data: dict) -> None:
    meta_path = paper_dir / "paper_metadata.json"
    _safe_write_json(meta_path, [data])
    print(f"[DEBUG] wrote metadata to {meta_path}")

    external_ids = data.get("externalIds", {}) or {}
    arxiv_id = external_ids.get("ArXiv")
    if arxiv_id:
        download_arxiv_tex(arxiv_id=arxiv_id, base_dir=paper_dir)

    sections_dir = paper_dir / "sections"
    sections_dir.mkdir(exist_ok=True)

    abstract = data.get("abstract")
    if abstract:
        _safe_write_text(sections_dir / "abstract.txt", abstract)

    tldr_obj = data.get("tldr")
    if isinstance(tldr_obj, dict) and tldr_obj.get("text"):
        _safe_write_text(sections_dir / "tldr.txt", tldr_obj["text"])

    semantic_id = data.get("paperId")
    if not semantic_id:
        print(f"[WARN] no semantic_id for {acl_id} → skip refs/cites.")
        return

    citation_count = data.get("citationCount", 0)
    reference_count = data.get("referenceCount", 0)

    ref_status, refs = get_paper_links(semantic_id, "references", reference_count)
    if ref_status == 200:
        _safe_write_json(paper_dir / "references_metadata.json", refs)

    cit_status, cits = get_paper_links(semantic_id, "citations", citation_count)
    if cit_status == 200:
        _safe_write_json(paper_dir / "citations_metadata.json", cits)

    if "ArXiv" not in external_ids:
        _safe_write_text(paper_dir / "no_arxiv.txt", "no arxiv for this paper")


def fetch_one_acl_id(paper: dict, base_dir: Path):
    acl_id = paper["id"]
    title = (paper.get("title") or "").strip()
    id_type = paper.get("id_type", "ACL")
    openreview_id = paper.get("openreview_id", "")
    input_pdf_url = paper.get("pdf_url", "")
    s2_key = os.getenv("SEMANTIC_SCHOLAR_API_KEY", "")
    print(
        f"[DEBUG] fetch_one_acl_id: id={acl_id} id_type={id_type} "
        f"title_len={len(title)} s2_key_present={'yes' if bool(s2_key) else 'no'} "
        f"s2_key_len={len(s2_key)}"
    )

    paper_dir = base_dir / acl_id
    ensure_dir(paper_dir)
    meta_path = paper_dir / "paper_metadata.json"

    if meta_path.exists():
        return

    status, data = 0, None
    fetch_label = f"{id_type}:{acl_id}"
    is_openreview = _treat_as_openreview(paper)
    openreview_meta = None
    attempted_title_search = False

    if is_openreview:
        try:
            openreview_meta = _query_openreview_for_paper(openreview_id or acl_id)
        except Exception as exc:
            print(f"[WARN] OpenReview lookup failed for {acl_id}: {exc}")
            openreview_meta = None

        if openreview_meta:
            _write_openreview_snapshot(paper_dir, openreview_meta)
            or_title = (openreview_meta.get("title") or title or "").strip()
            arxiv_id = (
                _best_arxiv_id(
                    openreview_meta.get("arxiv_id", ""),
                    openreview_meta.get("pdf_url", ""),
                    input_pdf_url,
                )
                or ""
            )
            if arxiv_id:
                print(f"[DEBUG] OpenReview -> found ArXiv {arxiv_id} for {acl_id}")
                status, data = get_paper(arxiv_id, id_type="ArXiv")
                fetch_label = f"ArXiv:{arxiv_id}"
                title = or_title or title
            elif or_title:
                print(f"[DEBUG] OpenReview -> no arXiv for {acl_id}, title-searching")
                status, data = _fetch_s2_by_title(or_title, acl_id)
                fetch_label = f"title:{or_title[:80]}"
                title = or_title
                attempted_title_search = True
            else:
                print(f"[WARN] OpenReview metadata for {acl_id} had neither title nor arXiv")
        else:
            print(f"[WARN] no OpenReview metadata for {acl_id} (openreview_id={openreview_id or acl_id})")

        if data is None and title and not attempted_title_search:
            print(f"[DEBUG] OpenReview fallback -> title-searching extracted title for {acl_id}")
            status, data = _fetch_s2_by_title(title, acl_id)
            fetch_label = f"title:{title[:80]}"
            attempted_title_search = True

    if data is None and not is_openreview:
        status, data = get_paper(acl_id, id_type=id_type)
        fetch_label = f"{id_type}:{acl_id}"

    if data is None and not attempted_title_search:
        print(
            f"[WARN] direct fetch failed for {fetch_label} "
            f"(status={status}) → trying title search with title_len={len(title)}"
        )
        status, data = _fetch_s2_by_title(title, acl_id)

    if status != 200 or data is None:
        print(f"[WARN] still no data for {acl_id} → skipping.")
        return

    _write_metadata_outputs(paper_dir, acl_id, data)
    print("[SUCCESS]")


def fetch_all_metadata(acl_ids_path: Path, out_dir: Path, start_from: str | None = None, resume: bool = False):
    raw = json.loads(acl_ids_path.read_text(encoding="utf-8"))
    papers = raw if isinstance(raw[0], dict) else [{"id": x, "title": ""} for x in raw]

    start_seen = start_from is None
    for paper in papers:
        pid = str(paper.get("id", ""))
        if not start_seen:
            if pid == start_from:
                start_seen = True
            else:
                continue
        if resume:
            paper_dir = out_dir / pid
            if (paper_dir / "paper_metadata.json").exists():
                continue
        fetch_one_acl_id(paper, out_dir)
    return "Meta Data Completed"


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--ids", type=str, required=True, help="Path to JSON file with paper IDs.")
    parser.add_argument("--outdir", type=str, default="papers", help="Output directory for metadata.")
    parser.add_argument("--start-from", type=str, default=None, help="Start from this paper ID.")
    parser.add_argument("--resume", action="store_true", help="Skip papers that already have paper_metadata.json.")
    args = parser.parse_args()

    ACL_IDS_PATH = Path(args.ids).expanduser().resolve()
    OUTDIR = Path(args.outdir).expanduser().resolve()

    if not ACL_IDS_PATH.exists():
        raise FileNotFoundError(f"Could not find {ACL_IDS_PATH}")

    print(f"[INFO] Using ID list from {ACL_IDS_PATH}")
    print(f"[INFO] Output will be saved to {OUTDIR}")

    start = time.time()
    fetch_all_metadata(acl_ids_path=ACL_IDS_PATH, out_dir=OUTDIR, start_from=args.start_from, resume=args.resume)
    print("done in", time.time() - start, "s")