#!/usr/bin/env python3
"""Fetch arXiv metadata and print BibTeX entries (stdlib only).

Usage:
    python3 arxiv2bib.py KEY=ARXIVID [KEY=ARXIVID ...]
    python3 arxiv2bib.py --list ids.txt        # one "KEY=ARXIVID" per line

Authors and titles come straight from the arXiv API, so they are the
authors' own metadata. Venues are NOT filled in here: arXiv only knows
about venues when authors add a comment/journal-ref, which is printed as a
`note` so you can confirm it and add `booktitle` by hand.
"""
import sys
import time
import urllib.request
import xml.etree.ElementTree as ET

API = "https://export.arxiv.org/api/query?id_list={ids}&max_results={n}"
NS = {"a": "http://www.w3.org/2005/Atom", "x": "http://arxiv.org/schemas/atom"}


def fetch(ids):
    url = API.format(ids=",".join(ids), n=len(ids))
    req = urllib.request.Request(url, headers={"User-Agent": "thesis-lit-tools/0.1"})
    for attempt in range(5):
        try:
            with urllib.request.urlopen(req, timeout=60) as r:
                return ET.fromstring(r.read())
        except Exception as e:  # rate limits / transient errors
            wait = 5 * (attempt + 1)
            print(f"% arXiv fetch failed ({e}); retrying in {wait}s", file=sys.stderr)
            time.sleep(wait)
    sys.exit("arXiv API unreachable")


def esc(s):
    return " ".join(s.split()).replace("&", r"\&")


def main(argv):
    pairs = []
    if argv[:1] == ["--list"]:
        lines = open(argv[1]).read().split()
        argv = [l for l in lines if "=" in l]
    for a in argv:
        key, aid = a.split("=", 1)
        pairs.append((key, aid))
    root = fetch([aid for _, aid in pairs])
    by_id = {}
    for e in root.findall("a:entry", NS):
        aid = e.find("a:id", NS).text.split("/abs/")[-1].split("v")[0]
        by_id[aid] = e
    for key, aid in pairs:
        e = by_id.get(aid)
        if e is None:
            print(f"% !! {key}: arXiv:{aid} NOT FOUND\n")
            continue
        title = esc(e.find("a:title", NS).text)
        authors = " and ".join(a.find("a:name", NS).text for a in e.findall("a:author", NS))
        year = e.find("a:published", NS).text[:4]
        notes = [x.text for x in (e.find("x:comment", NS), e.find("x:journal_ref", NS)) if x is not None]
        print(f"@misc{{{key},")
        print(f"  title         = {{{{{title}}}}},")
        print(f"  author        = {{{authors}}},")
        print(f"  year          = {{{year}}},")
        print(f"  eprint        = {{{aid}}},")
        print(f"  archivePrefix = {{arXiv}},")
        if notes:
            print(f"  note          = {{arXiv comment: {esc(' | '.join(notes))}}},")
        print("}\n")


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