#!/usr/bin/env python3
"""papers.py - find, resolve, download and cite scientific papers (stdlib only).

    python papers.py search "boolean operations on triangle meshes" --max 25 [--since 2018]
    python papers.py resolve 1706.03762             # arXiv id, DOI or "title": find the peer-reviewed version
    python papers.py bibtex 10.1145/3306346.3322959 # DOI / arXiv id / dblp key -> BibTeX entry
    python papers.py download 1706.03762 --out papers/   # PDF (arXiv or open-access copy)
    python papers.py text papers/vaswani2017attention.pdf [--pages 12]   # PDF -> .txt (pymupdf)

search  merges arXiv, OpenAlex and Semantic Scholar; prints a ranked table and writes candidates.json.
resolve prints, for one paper, the best version: journal/conference (DOI, venue, year) if one
        exists, otherwise the arXiv preprint - plus the BibTeX key you should use.
bibtex  prints a cleaned BibTeX entry with a stable key  <lastname><year><firstword>.
All network calls have a 20 s timeout and identify themselves; be polite (<= 1 request/s).
Sources: export.arxiv.org (Atom), api.openalex.org, api.semanticscholar.org, api.crossref.org,
dblp.org, doi.org (content negotiation). Set CONTACT_EMAIL for the Crossref/OpenAlex polite pools.
"""
import json, os, re, sys, time, unicodedata, urllib.parse, urllib.request
import xml.etree.ElementTree as ET
from pathlib import Path

UA = f"papers.py/1.0 (mailto:{os.environ.get('CONTACT_EMAIL', 'research@example.org')})"
ARXIV_RE = re.compile(r"(?:arxiv[:/]|abs/|pdf/)?(\d{4}\.\d{4,5})(?:v\d+)?", re.I)
DOI_RE = re.compile(r"\b(10\.\d{4,9}/[^\s\"<>]+)", re.I)


# ----------------------------------------------------------------- helpers
def get(url, accept="application/json", retries=2):
    req = urllib.request.Request(url, headers={"User-Agent": UA, "Accept": accept})
    for attempt in range(retries + 1):
        try:
            with urllib.request.urlopen(req, timeout=20) as r:
                return r.read().decode("utf-8", "replace")
        except Exception as e:  # noqa: BLE001
            if attempt == retries:
                print(f"  ! {url[:80]}...: {e}", file=sys.stderr)
                return None
            time.sleep(2 * (attempt + 1))


def norm_title(t):
    t = unicodedata.normalize("NFKD", t or "").encode("ascii", "ignore").decode()
    return re.sub(r"[^a-z0-9]+", " ", t.lower()).strip()


def bib_key(authors, year, title):
    last = (authors[0].split()[-1] if authors else "anon")
    last = re.sub(r"[^a-z]", "", unicodedata.normalize("NFKD", last).encode("ascii", "ignore").decode().lower()) or "anon"
    stop = {"a", "an", "the", "on", "of", "for", "and", "to", "in", "with", "towards", "toward"}
    first = next((w for w in re.findall(r"[a-z0-9]+", (title or "").lower()) if w not in stop), "paper")
    return f"{last}{year or 'nd'}{first}"


def is_arxiv_doi(doi):
    return bool(doi) and doi.lower().startswith("10.48550/arxiv")


# ----------------------------------------------------------------- parsers (pure: testable offline)
def parse_arxiv(atom):
    ns = {"a": "http://www.w3.org/2005/Atom", "x": "http://arxiv.org/schemas/atom"}
    out = []
    for e in ET.fromstring(atom).findall("a:entry", ns):
        aid = ARXIV_RE.search(e.findtext("a:id", "", ns))
        out.append({
            "arxiv": aid.group(1) if aid else None,
            "title": " ".join(e.findtext("a:title", "", ns).split()),
            "authors": [a.findtext("a:name", "", ns) for a in e.findall("a:author", ns)],
            "year": (e.findtext("a:published", "", ns) or "")[:4],
            "abstract": " ".join(e.findtext("a:summary", "", ns).split()),
            "doi": (e.findtext("x:doi", None, ns) or None),
            "venue": (e.findtext("x:journal_ref", None, ns) or None),   # authors' own note, e.g. "NeurIPS 2017"
            "pdf": f"https://arxiv.org/pdf/{aid.group(1)}" if aid else None,
            "source": "arxiv",
        })
    return out


def parse_openalex(js):
    out = []
    for w in json.loads(js).get("results", []):
        ids = w.get("ids", {}) or {}
        doi = (w.get("doi") or "").replace("https://doi.org/", "") or None
        loc = w.get("primary_location") or {}
        src = loc.get("source") or {}
        best = w.get("best_oa_location") or {}
        # abstract comes as an inverted index
        inv = w.get("abstract_inverted_index") or {}
        words = sorted((p, t) for t, ps in inv.items() for p in ps)
        out.append({
            "arxiv": next((m.group(1) for u in (best.get("pdf_url"), best.get("landing_page_url"))
                           if u and "arxiv.org" in u and (m := ARXIV_RE.search(u))), None),
            "title": w.get("title") or "",
            "authors": [a.get("author", {}).get("display_name", "") for a in w.get("authorships", [])],
            "year": str(w.get("publication_year") or ""),
            "abstract": " ".join(t for _, t in words),
            "doi": doi,
            "venue": src.get("display_name"),
            "type": w.get("type"),                       # article, preprint, book-chapter, ...
            "citations": w.get("cited_by_count", 0),
            "pdf": best.get("pdf_url"),
            "source": "openalex",
        })
    return out


def parse_s2(js):
    data = json.loads(js)
    items = data.get("data", [data]) if isinstance(data, dict) else data
    out = []
    for p in items:
        if not p:
            continue
        ext = p.get("externalIds") or {}
        pv = p.get("publicationVenue") or {}
        out.append({
            "arxiv": ext.get("ArXiv"),
            "title": p.get("title") or "",
            "authors": [a.get("name", "") for a in p.get("authors", [])],
            "year": str(p.get("year") or ""),
            "abstract": p.get("abstract") or "",
            "doi": ext.get("DOI"),
            "venue": p.get("venue") or pv.get("name"),
            "venue_type": pv.get("type"),                # journal / conference
            "types": p.get("publicationTypes") or [],    # JournalArticle, Conference, ...
            "citations": p.get("citationCount", 0),
            "pdf": (p.get("openAccessPdf") or {}).get("url"),
            "source": "s2",
        })
    return out


def parse_crossref(js, want_title):
    """Best Crossref item whose title matches (journal/proceedings only), else None."""
    want = norm_title(want_title)
    for it in json.loads(js).get("message", {}).get("items", []):
        title = (it.get("title") or [""])[0]
        if norm_title(title) != want or it.get("type") not in ("journal-article", "proceedings-article", "book-chapter"):
            continue
        date = (it.get("published-print") or it.get("published-online") or it.get("issued") or {}).get("date-parts", [[None]])[0][0]
        return {
            "doi": it.get("DOI"), "title": title, "year": str(date or ""),
            "venue": (it.get("container-title") or [None])[0] or (it.get("event") or {}).get("name"),
            "type": it.get("type"),
            "authors": [f"{a.get('given', '')} {a.get('family', '')}".strip() for a in it.get("author", [])],
            "source": "crossref",
        }
    return None


def parse_dblp(js, want_title):
    """dblp hit for the title, preferring the conference/journal record over the CoRR (arXiv) one."""
    want = norm_title(want_title)
    hits = [h.get("info", {}) for h in json.loads(js).get("result", {}).get("hits", {}).get("hit", [])]
    hits = [h for h in hits if norm_title(h.get("title", "").rstrip(".")) == want]
    hits.sort(key=lambda h: (h.get("venue") == "CoRR", h.get("type") not in ("Conference and Workshop Papers", "Journal Articles")))
    if not hits:
        return None
    h = hits[0]
    return {"key": h.get("key"), "venue": h.get("venue"), "year": h.get("year"), "type": h.get("type"),
            "doi": h.get("doi"), "bib": f"https://dblp.org/rec/{h.get('key')}.bib", "source": "dblp"}


def merge(cands):
    """Deduplicate by DOI / arXiv id / normalized title; keep the most informative record."""
    seen, out = {}, []
    for c in cands:
        keys = [k for k in (c.get("doi") and c["doi"].lower(), c.get("arxiv"), norm_title(c["title"])) if k]
        hit = next((seen[k] for k in keys if k in seen), None)
        if hit is None:
            hit = dict(c); out.append(hit)
        else:
            for k, v in c.items():
                if v and not hit.get(k):
                    hit[k] = v
            hit["citations"] = max(hit.get("citations") or 0, c.get("citations") or 0)
        for k in keys:
            seen[k] = hit
    return out


def arxiv_bibtex(rec):
    key = bib_key(rec["authors"], rec["year"], rec["title"])
    return (f"@misc{{{key},\n  title = {{{rec['title']}}},\n  author = {{{' and '.join(rec['authors'])}}},\n"
            f"  year = {{{rec['year']}}},\n  eprint = {{{rec['arxiv']}}},\n  archivePrefix = {{arXiv}},\n"
            f"  note = {{arXiv preprint}},\n  url = {{https://arxiv.org/abs/{rec['arxiv']}}}\n}}")


def rekey_bibtex(entry, key):
    """Replace the entry key (doi.org / dblp keys are ugly) and tidy whitespace."""
    entry = re.sub(r"^\s*@(\w+)\s*\{\s*[^,]*,", lambda m: f"@{m.group(1).lower()}{{{key},", entry.strip(), count=1)
    return re.sub(r"\n\s*\n", "\n", entry)


# ----------------------------------------------------------------- network commands
def search(query, max_n=25, since=None):
    q = urllib.parse.quote(query)
    cands = []
    atom = get(f"https://export.arxiv.org/api/query?search_query=all:{q}&max_results={max_n}&sortBy=relevance", "application/atom+xml")
    if atom:
        cands += parse_arxiv(atom)
    oa = get(f"https://api.openalex.org/works?search={q}&per-page={max_n}&mailto={urllib.parse.quote(UA.split('mailto:')[1][:-1])}")
    if oa:
        cands += parse_openalex(oa)
    time.sleep(1)
    s2 = get(f"https://api.semanticscholar.org/graph/v1/paper/search?query={q}&limit={max_n}"
             "&fields=title,authors,year,abstract,venue,publicationVenue,externalIds,citationCount,openAccessPdf,publicationTypes")
    if s2:
        cands += parse_s2(s2)
    merged = [c for c in merge(cands) if not since or (c.get("year") or "0") >= str(since)]
    merged.sort(key=lambda c: -(c.get("citations") or 0))
    for c in merged:
        c["key"] = bib_key(c.get("authors", []), c.get("year"), c.get("title"))
    Path("candidates.json").write_text(json.dumps(merged, indent=2, ensure_ascii=False), encoding="utf-8")
    print(f"{len(merged)} candidates (candidates.json). cites  year  venue                      key / title")
    for c in merged:
        venue = (c.get("venue") or ("arXiv" if c.get("arxiv") else "?"))[:26]
        print(f"{c.get('citations') or 0:6d}  {c.get('year') or '    '}  {venue:26s} {c['key']}  {c['title'][:70]}")
    return merged


def resolve(ident):
    """Find the peer-reviewed version of a paper given an arXiv id, DOI or title."""
    arxiv = ARXIV_RE.fullmatch(ident.strip()) or (ARXIV_RE.search(ident) if "arxiv" in ident.lower() else None)
    doi = DOI_RE.search(ident)
    rec = None
    if arxiv or (doi and not is_arxiv_doi(doi.group(1))):
        pid = f"arXiv:{arxiv.group(1)}" if arxiv else doi.group(1)
        js = get(f"https://api.semanticscholar.org/graph/v1/paper/{pid}?fields=title,authors,year,venue,publicationVenue,externalIds,citationCount,openAccessPdf,publicationTypes")
        rec = parse_s2(js)[0] if js and not js.startswith('{"error"') else None
    title = rec["title"] if rec else ident.strip()
    result = {"query": ident, "title": title, "preprint": None, "published": None}
    if rec and rec.get("arxiv"):
        result["preprint"] = {"arxiv": rec["arxiv"], "url": f"https://arxiv.org/abs/{rec['arxiv']}", "pdf": f"https://arxiv.org/pdf/{rec['arxiv']}"}
    # 1. Semantic Scholar already knows a non-arXiv DOI + venue?
    if rec and rec.get("doi") and not is_arxiv_doi(rec["doi"]) and rec.get("venue"):
        result["published"] = {"doi": rec["doi"], "venue": rec["venue"], "year": rec["year"], "via": "semanticscholar",
                               "type": rec.get("venue_type") or ",".join(rec.get("types", []))}
    time.sleep(1)
    # 2. Crossref by exact title (journal/proceedings only)
    if not result["published"]:
        cr = get(f"https://api.crossref.org/works?query.bibliographic={urllib.parse.quote(title)}&rows=5&mailto={urllib.parse.quote(UA.split('mailto:')[1][:-1])}")
        hit = parse_crossref(cr, title) if cr else None
        if hit:
            result["published"] = {"doi": hit["doi"], "venue": hit["venue"], "year": hit["year"], "via": "crossref", "type": hit["type"]}
    # 3. dblp (CS): venue + a clean .bib
    db = get(f"https://dblp.org/search/publ/api?q={urllib.parse.quote(title)}&format=json&h=10")
    hit = parse_dblp(db, title) if db else None
    if hit and hit["venue"] != "CoRR":
        pub = result["published"] or {}
        pub.update({k: v for k, v in {"venue": hit["venue"], "year": hit["year"], "doi": pub.get("doi") or hit.get("doi"),
                                      "dblp": hit["key"], "bib": hit["bib"], "type": hit["type"]}.items() if v})
        pub["via"] = pub.get("via") or "dblp"
        result["published"] = pub
    authors = rec["authors"] if rec else []
    result["key"] = bib_key(authors, (result["published"] or {}).get("year") or (rec or {}).get("year"), title)
    print(json.dumps(result, indent=2, ensure_ascii=False))
    if result["published"]:
        print(f"\n-> cite the published version: {result['published'].get('venue')} ({result['published'].get('year')}), "
              f"DOI {result['published'].get('doi')}; bibtex: python papers.py bibtex {result['published'].get('doi') or result['published'].get('dblp')}")
    elif result["preprint"]:
        print(f"\n-> no peer-reviewed version found: cite as arXiv preprint {result['preprint']['arxiv']} and say so in research.md")
    return result


def bibtex(ident):
    ident = ident.strip()
    if ident.startswith("conf/") or ident.startswith("journals/"):                      # dblp key
        entry = get(f"https://dblp.org/rec/{ident}.bib", "text/plain")
        rec = None
    elif (m := DOI_RE.search(ident)) and not is_arxiv_doi(m.group(1)):
        entry = get(f"https://doi.org/{m.group(1)}", "application/x-bibtex")
        rec = None
    else:
        aid = ARXIV_RE.search(ident)
        if not aid:
            print("give a DOI, an arXiv id or a dblp key", file=sys.stderr); return None
        atom = get(f"https://export.arxiv.org/api/query?id_list={aid.group(1)}", "application/atom+xml")
        rec = parse_arxiv(atom)[0] if atom else None
        entry = arxiv_bibtex(rec) if rec else None
    if not entry:
        print("no BibTeX found", file=sys.stderr); return None
    m_auth = re.search(r"author\s*=\s*[{\"](.+?)[}\"]\s*,?\s*\n", entry, re.S | re.I)
    m_year = re.search(r"year\s*=\s*[{\"]?(\d{4})", entry, re.I)
    m_title = re.search(r"title\s*=\s*[{\"](.+?)[}\"]\s*,?\s*\n", entry, re.S | re.I)
    authors = [a.strip() for a in (m_auth.group(1) if m_auth else "").replace("\n", " ").split(" and ")]
    authors = [(" ".join(reversed(a.split(", "))) if ", " in a else a) for a in authors]
    key = bib_key(authors, m_year.group(1) if m_year else "", re.sub(r"[{}]", "", m_title.group(1)) if m_title else "")
    entry = rekey_bibtex(entry, key)
    print(entry)
    return entry


def download(ident, out="papers"):
    Path(out).mkdir(parents=True, exist_ok=True)
    aid = ARXIV_RE.search(ident)
    url, name = None, None
    if ident.startswith("http"):
        url, name = ident, re.sub(r"[^\w.-]", "_", ident.rsplit("/", 1)[-1]) or "paper.pdf"
    elif aid:
        url, name = f"https://arxiv.org/pdf/{aid.group(1)}", f"arxiv_{aid.group(1).replace('.', '_')}.pdf"
    elif (m := DOI_RE.search(ident)):
        oa = get(f"https://api.openalex.org/works/https://doi.org/{m.group(1)}")
        best = (json.loads(oa).get("best_oa_location") or {}) if oa else {}
        url, name = best.get("pdf_url"), re.sub(r"[^\w.-]", "_", m.group(1)) + ".pdf"
        if not url:
            print("no open-access PDF known for this DOI (try the arXiv version, or the publisher page via the library)"); return None
    if not url:
        print("give an arXiv id, a DOI or a URL", file=sys.stderr); return None
    req = urllib.request.Request(url, headers={"User-Agent": UA, "Accept": "application/pdf"})
    try:
        with urllib.request.urlopen(req, timeout=60) as r:
            data = r.read()
    except Exception as e:  # noqa: BLE001
        print(f"download failed: {e}", file=sys.stderr); return None
    if not data.startswith(b"%PDF"):
        print(f"not a PDF (got {data[:40]!r}) - probably a landing page or a captcha"); return None
    path = Path(out) / name
    path.write_bytes(data)
    print(f"saved {path} ({len(data) // 1024} KB)")
    return path


def text(pdf, pages=None):
    try:
        import pymupdf as fitz
    except ImportError:
        try:
            import fitz
        except ImportError:
            print("pip install pymupdf", file=sys.stderr); return None
    doc = fitz.open(pdf)
    n = min(doc.page_count, pages or doc.page_count)
    parts = [f"\n\n===== page {i + 1} =====\n" + doc[i].get_text() for i in range(n)]
    out = Path(pdf).with_suffix(".txt")
    out.write_text("".join(parts), encoding="utf-8")
    print(f"{out}: {n}/{doc.page_count} pages, {sum(len(p) for p in parts)} chars")
    return out


def main(argv):
    if not argv or argv[0] in ("-h", "--help"):
        print(__doc__); return 0
    cmd, args = argv[0], argv[1:]
    opt = {a.split("=")[0]: True for a in args if a.startswith("--")}
    def val(flag, default=None):
        if flag in args:
            i = args.index(flag); return args[i + 1]
        return default
    pos = [a for i, a in enumerate(args) if not a.startswith("--") and (i == 0 or not args[i - 1].startswith("--"))]
    if cmd == "search":
        search(" ".join(pos), int(val("--max", 25)), val("--since"))
    elif cmd == "resolve":
        resolve(" ".join(pos))
    elif cmd == "bibtex":
        bibtex(" ".join(pos))
    elif cmd == "download":
        download(pos[0], val("--out", "papers"))
    elif cmd == "text":
        text(pos[0], int(val("--pages")) if val("--pages") else None)
    else:
        print(__doc__); return 2
    return 0


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