"""Download the raw inputs: NASA Exoplanet Archive parameters and arXiv metadata for every citation.

Raw files go to RAW_DIR (default fermi_raw/mdwarf_early).  A small snapshot of each
is also written to data/ so the numbers in report.md can be regenerated if a service is down
(run_all.py --offline) and so that a referee can diff against a fresh download.
"""
import csv
import io
import json
import os
import re
import shutil
import sys
import time
import urllib.parse
import urllib.request
import xml.etree.ElementTree as ET

import pandas as pd

from common import DATA_DIR, RAW_DIR
from curated_outcomes import OUTCOMES, EXTRA_REFS, PREDICTION_TARGETS

NEA_TAP = "https://exoplanetarchive.ipac.caltech.edu/TAP/sync"
NEA_COLS = ("pl_name,hostname,st_spectype,st_teff,st_mass,st_rad,st_lum,st_age,"
            "pl_rade,pl_radeerr1,pl_radeerr2,pl_bmasse,pl_bmasseerr1,pl_bmasseerr2,pl_bmasselim,"
            "pl_bmassprov,pl_insol,pl_eqt,pl_orbper,pl_orbsmax,pl_dens,sy_dist,disc_year")


def _get(url, tries=4, timeout=120):
    last = None
    for k in range(tries):
        try:
            with urllib.request.urlopen(url, timeout=timeout) as r:
                return r.read().decode()
        except Exception as e:  # noqa: BLE001
            last = e
            time.sleep(4 * (k + 1))
    raise RuntimeError(f"download failed: {url[:120]}... ({last})")


def fetch_nea(refresh=False):
    """Parameters for (i) every curated planet and prediction target, (ii) all small temperate M-dwarf planets."""
    raw1 = os.path.join(RAW_DIR, "nea_pscomppars_selected.csv")
    raw2 = os.path.join(RAW_DIR, "nea_pscomppars_temperate_mdwarf.csv")
    if refresh or not os.path.exists(raw1):
        names = sorted({o["planet"] for o in OUTCOMES} | set(PREDICTION_TARGETS))
        inl = ",".join("'" + n.replace("'", "''") + "'" for n in names)
        q = f"select {NEA_COLS} from pscomppars where pl_name in ({inl})"
        txt = _get(NEA_TAP + "?" + urllib.parse.urlencode(dict(query=q, format="csv")))
        open(raw1, "w").write(txt)
    if refresh or not os.path.exists(raw2):
        q = (f"select {NEA_COLS} from pscomppars where st_teff < 3900 and pl_rade < 1.8 "
             "and st_lum is not null and pl_orbsmax is not null")
        txt = _get(NEA_TAP + "?" + urllib.parse.urlencode(dict(query=q, format="csv")))
        open(raw2, "w").write(txt)
    shutil.copy(raw1, os.path.join(DATA_DIR, "nea_snapshot_selected.csv"))
    shutil.copy(raw2, os.path.join(DATA_DIR, "nea_snapshot_small_mdwarf_planets.csv"))
    return pd.read_csv(raw1), pd.read_csv(raw2)


def all_arxiv_ids():
    ids = []
    for o in OUTCOMES:
        ids += o["refs"].split(";")
    ids += EXTRA_REFS.split(";")
    for o in OUTCOMES:  # ids quoted inside free-text XUV notes
        ids += re.findall(r"arXiv:(\d{4}\.\d{4,5})", o.get("xuv", ""))
    out = []
    for i in ids:
        i = i.strip()
        if i and i not in out:
            out.append(i)
    return out


def fetch_arxiv(refresh=False):
    raw = os.path.join(RAW_DIR, "arxiv_metadata.json")
    meta = {}
    if os.path.exists(raw) and not refresh:
        meta = json.load(open(raw))
    todo = [i for i in all_arxiv_ids() if i not in meta]
    ns = {"a": "http://www.w3.org/2005/Atom", "x": "http://arxiv.org/schemas/atom"}
    for k in range(0, len(todo), 20):
        chunk = todo[k:k + 20]
        url = "https://export.arxiv.org/api/query?" + urllib.parse.urlencode(
            dict(id_list=",".join(chunk), max_results=50))
        root = ET.fromstring(_get(url))
        for e in root.findall("a:entry", ns):
            id_ = re.sub(r"v\d+$", "", e.find("a:id", ns).text.split("/abs/")[-1])
            jr = e.find("x:journal_ref", ns)
            authors = [x.find("a:name", ns).text for x in e.findall("a:author", ns)]
            meta[id_] = dict(
                title=re.sub(r"\s+", " ", e.find("a:title", ns).text).strip(),
                published=e.find("a:published", ns).text[:10],
                first_author=authors[0], n_authors=len(authors),
                journal_ref=jr.text if jr is not None else "",
                abstract=re.sub(r"\s+", " ", e.find("a:summary", ns).text).strip())
        time.sleep(3)
    json.dump(meta, open(raw, "w"), indent=1)
    rows = []
    for i in all_arxiv_ids():
        m = meta.get(i)
        if m is None:
            rows.append(dict(arxiv_id=i, first_author="NOT FOUND", n_authors="", published="",
                             title="", journal_ref="", url=f"https://arxiv.org/abs/{i}"))
        else:
            rows.append(dict(arxiv_id=i, first_author=m["first_author"], n_authors=m["n_authors"],
                             published=m["published"], title=m["title"],
                             journal_ref=m["journal_ref"], url=f"https://arxiv.org/abs/{i}"))
    pd.DataFrame(rows).to_csv(os.path.join(DATA_DIR, "citations_verified.csv"), index=False,
                              quoting=csv.QUOTE_MINIMAL)
    return meta


def load_offline():
    """Use the snapshots in data/ (no network)."""
    sel = pd.read_csv(os.path.join(DATA_DIR, "nea_snapshot_selected.csv"))
    temp = pd.read_csv(os.path.join(DATA_DIR, "nea_snapshot_small_mdwarf_planets.csv"))
    cit = pd.read_csv(os.path.join(DATA_DIR, "citations_verified.csv"), dtype=str).fillna("")
    meta = {r.arxiv_id: dict(title=r.title, published=r.published, first_author=r.first_author,
                             n_authors=r.n_authors, journal_ref=r.journal_ref)
            for r in cit.itertuples()}
    return sel, temp, meta


def main(refresh=False, offline=False):
    if offline:
        return load_offline()
    sel, temp = fetch_nea(refresh)
    meta = fetch_arxiv(refresh)
    return sel, temp, meta


if __name__ == "__main__":
    s, t, m = main(refresh="--refresh" in sys.argv, offline="--offline" in sys.argv)
    print(len(s), "selected planets;", len(t), "small M-dwarf planets;", len(m), "arXiv records")