#!/usr/bin/env python3
"""Keyword search over OpenAlex and arXiv for the literature review (stdlib only).

Usage:
    python3 search.py openalex 'TERMS' --out NAME [--since 2020-01-01] [--limit 200]
    python3 search.py arxiv    'ARXIV_QUERY' --out NAME [--since 2020-01-01] [--limit 200]

OpenAlex TERMS: '|'-separated phrases (OR), ' AND ' between concept blocks, e.g.
    '"qk-norm"|"query-key normalization"'
arXiv ARXIV_QUERY: arXiv API syntax, e.g.
    'abs:"qk norm" OR abs:"query-key normalization"'

Writes ../search/results/NAME.csv and appends a row to ../search/search-log.csv
(same format as s2.py). Raw responses are cached in ../search/.cache/.
"""
import csv
import datetime as dt
import hashlib
import json
import os
import sys
import time
import urllib.parse
import urllib.request
import xml.etree.ElementTree as ET

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")
COLS = ["source_id", "arxiv", "doi", "year", "title", "first_author", "n_authors", "venue", "citations", "abstract"]
NS = {"a": "http://www.w3.org/2005/Atom"}


def fetch(url, raw=False):
    key = os.path.join(CACHE, hashlib.sha1(url.encode()).hexdigest())
    if os.path.exists(key):
        body = open(key, "rb").read()
    else:
        for attempt in range(6):
            try:
                req = urllib.request.Request(url, headers={"User-Agent": "msc-thesis-lit-search/0.1"})
                with urllib.request.urlopen(req, timeout=60) as r:
                    body = r.read()
                break
            except Exception as e:
                wait = 5 * 2 ** attempt
                print(f"  {e}; retry in {wait}s", file=sys.stderr)
                time.sleep(wait)
        else:
            sys.exit(f"gave up on {url}")
        os.makedirs(CACHE, exist_ok=True)
        open(key, "wb").write(body)
        time.sleep(1)
    return body if raw else json.loads(body)


def clean(s):
    return " ".join((s or "").split())


def openalex_abstract(inv):
    if not inv:
        return ""
    pos = sorted((i, w) for w, idx in inv.items() for i in idx)
    return " ".join(w for _, w in pos)


def openalex(terms, since, limit):
    # ' AND ' separates concept blocks; each block is an OR-list of phrases.
    blocks = [f"title_and_abstract.search:{b.strip()}" for b in terms.split(" AND ")]
    filt = ",".join(blocks + ([f"from_publication_date:{since}"] if since else []))
    rows, page, total = [], 1, 0
    while len(rows) < limit:
        url = "https://api.openalex.org/works?" + urllib.parse.urlencode(
            {"filter": filt, "per_page": 100, "page": page}
        )
        d = fetch(url)
        total = d["meta"]["count"]
        for w in d["results"]:
            ids = w.get("ids") or {}
            doi = (w.get("doi") or "").replace("https://doi.org/", "")
            arxiv = doi.lower().split("arxiv.")[-1] if "arxiv." in doi.lower() else ""
            src = ((w.get("primary_location") or {}).get("source") or {}).get("display_name", "")
            auth = w.get("authorships") or []
            rows.append({
                "source_id": ids.get("openalex", "").split("/")[-1],
                "arxiv": arxiv, "doi": doi, "year": w.get("publication_year") or "",
                "title": clean(w.get("title")),
                "first_author": (auth[0]["author"]["display_name"] if auth else ""),
                "n_authors": len(auth), "venue": src, "citations": w.get("cited_by_count") or 0,
                "abstract": clean(openalex_abstract(w.get("abstract_inverted_index"))),
            })
        if len(d["results"]) < 100:
            break
        page += 1
    return rows[:limit], total


def arxiv(query, since, limit):
    rows, start, total = [], 0, 0
    while len(rows) < limit:
        url = "https://export.arxiv.org/api/query?" + urllib.parse.urlencode(
            {"search_query": query, "start": start, "max_results": 100, "sortBy": "relevance"}
        )
        root = ET.fromstring(fetch(url, raw=True))
        total = int(root.find("{http://a9.com/-/spec/opensearch/1.1/}totalResults").text)
        entries = root.findall("a:entry", NS)
        for e in entries:
            year = e.find("a:published", NS).text[:10]
            if since and year < since:
                continue
            aid = e.find("a:id", NS).text.split("/abs/")[-1].rsplit("v", 1)[0]
            au = e.findall("a:author", NS)
            rows.append({
                "source_id": aid, "arxiv": aid, "doi": "", "year": year[:4],
                "title": clean(e.find("a:title", NS).text),
                "first_author": au[0].find("a:name", NS).text if au else "",
                "n_authors": len(au), "venue": "arXiv", "citations": "",
                "abstract": clean(e.find("a:summary", NS).text),
            })
        if len(entries) < 100:
            break
        start += 100
        time.sleep(3)  # arXiv asks for >=3s between calls
    return rows[:limit], total


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

    src, query, out = a[0], a[1], opt("--out")
    since, limit = opt("--since", "2020-01-01"), int(opt("--limit", 200))
    rows, total = {"openalex": openalex, "arxiv": arxiv}[src](query, since, limit)
    os.makedirs(RESULTS, exist_ok=True)
    path = os.path.join(RESULTS, out + ".csv")
    with open(path, "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=COLS)
        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(), {"openalex": "OpenAlex API", "arxiv": "arXiv API"}[src],
                    "search", query, f"since={since}", total, len(rows), os.path.relpath(path, SEARCH_DIR)])
    print(f"{src} {out}: {total} total, {len(rows)} saved")


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