"""Download every external dataset used by the study into the raw directory.

Sources (all public, plain HTTPS):
  * SETI Institute Technosearch (radio / optical / archival lists)  -> parsed to data/technosearch_*.csv
  * VizieR TAP: Wlodarczyk-Sroka+ 2020 (J/MNRAS/498/5720), Harp+ 2016 (J/AJ/152/181),
                Isaacson+ 2017 (J/PASP/129/E4501), Gaia Catalogue of Nearby Stars (J/A+A/649/A6),
                Baumgardt & Hilker 2018 globular-cluster masses (J/MNRAS/478/1520)
  * CDS XMatch: Harp+ 2016 targets x Gaia DR3 (I/355/gaiadr3) for parallaxes
  * Gaia archive TAP: binned M_G - luminosity relation from DR3 FLAME (for the Dyson-swarm sensitivity)

Files are cached: an existing non-empty file is not downloaded again (delete it to refresh).
"""
import csv
import html
import io
import os
import re
import sys
import time
import urllib.parse
import urllib.request
import uuid

from common import RAW, DATA

VIZ = os.path.join(RAW, "vizier")
TAP = "https://tapvizier.cds.unistra.fr/TAPVizieR/tap/sync"
XMATCH = "http://cdsxmatch.u-strasbg.fr/xmatch/api/v1/sync"
# technosearch.seti.org answers 403 to the default Python user agent
UA = "Mozilla/5.0 (X11; Linux x86_64) seti_meta-study/1.0 (research script; urllib)"


def _get(url, data=None, headers=None, timeout=600, tries=4, pause=5):
    last = None
    for i in range(tries):
        try:
            hdr = {"User-Agent": UA}
            hdr.update(headers or {})
            req = urllib.request.Request(url, data=data, headers=hdr)
            with urllib.request.urlopen(req, timeout=timeout) as r:
                return r.read()
        except Exception as e:  # noqa
            last = e
            time.sleep(pause * (i + 1))
    raise RuntimeError(f"download failed: {url}: {last}")


def cached(path):
    return os.path.exists(path) and os.path.getsize(path) > 200


def tap(query, out, maxrec=500000):
    if cached(out):
        return
    params = urllib.parse.urlencode({"REQUEST": "doQuery", "LANG": "ADQL", "FORMAT": "csv",
                                     "MAXREC": str(maxrec), "QUERY": query})
    print("  TAP ->", os.path.basename(out)); sys.stdout.flush()
    b = _get(TAP + "?" + params)
    if b.lstrip().startswith(b"<?xml"):
        raise RuntimeError("TAP error: " + b[:600].decode(errors="replace"))
    with open(out, "wb") as f:
        f.write(b)


def fetch_technosearch():
    for name in ("radio-list", "optical-list", "archival-list"):
        out = os.path.join(RAW, f"technosearch_{name}.html")
        if not cached(out):
            print("  Technosearch ->", name)
            with open(out, "wb") as f:
                f.write(_get(f"https://technosearch.seti.org/{name}/", timeout=90))
        s = open(out, encoding="utf8").read()
        rows = []
        for t in re.findall(r"<table.*?</table>", s, flags=re.S):
            for r in re.findall(r"<tr[^>]*>(.*?)</tr>", t, flags=re.S):
                cells = re.findall(r"<t[hd][^>]*>(.*?)</t[hd]>", r, flags=re.S)
                links = [l for l in re.findall(r'href="([^"]+)"', r)]
                clean = []
                for c in cells:
                    c = re.sub(r"<sup>(.*?)</sup>", r"^\1", c, flags=re.S)
                    c = re.sub(r"<[^>]+>", "", c)
                    clean.append(re.sub(r"\s+", " ", html.unescape(c)).strip())
                links = ["https://technosearch.seti.org" + l if l.startswith("/") else l for l in links]
                rows.append(clean + [" | ".join(links)])
        if rows:
            rows[0][-1] = "links"
        with open(os.path.join(DATA, f"technosearch_{name.replace('-list','')}.csv"), "w", newline="", encoding="utf8") as f:
            csv.writer(f).writerows(rows)
        print(f"  technosearch {name}: {len(rows)-1} rows")


def fetch_vizier():
    tap('SELECT recno,RA_ICRS,DE_ICRS,Plx,RPlx,Gmag,Teff,Target,RATdeg,DETdeg,PlxT,rest,InstT,InstS,FWHM,'
        '"Offset-arcmin",TelResGaus,"min-detec-EIRP","Price-target?" FROM "J/MNRAS/498/5720/catalog"',
        os.path.join(VIZ, "ws2020_catalog.csv"))
    tap('SELECT recno,S,Name,Percent,BW,SimbadName,"_RA","_DE" FROM "J/AJ/152/181/table7"',
        os.path.join(VIZ, "harp2016_table7.csv"))
    tap('SELECT * FROM "J/AJ/152/181/table6"', os.path.join(VIZ, "harp2016_table6.csv"))
    tap('SELECT * FROM "J/PASP/129/E4501/table1"', os.path.join(VIZ, "isaacson2017_table1.csv"))
    tap('SELECT GaiaEDR3,RA_ICRS,DE_ICRS,Plx,"Gmag",BPmag,RPmag,GCNSprob,WDprob,Dist50,W3mag,W4mag '
        'FROM "J/A+A/649/A6/table1c"', os.path.join(VIZ, "gcns_table1c.csv"))
    tap('SELECT Cluster,Dist,Mass,e_Mass FROM "J/MNRAS/478/1520/table2"', os.path.join(VIZ, "baumgardt2018_gc.csv"))


def xmatch_harp():
    out = os.path.join(VIZ, "harp2016_x_gaiadr3.csv")
    if cached(out):
        return
    rows = []
    for fn, tag in (("harp2016_table7.csv", "t7"), ("harp2016_table6.csv", "t6")):
        for r in csv.DictReader(open(os.path.join(VIZ, fn), encoding="utf8")):
            if r["_RA"] and r["_DE"]:
                rows.append((f"{tag}_{r['recno']}", r["_RA"], r["_DE"]))
    buf = "id,ra,dec\n" + "\n".join(",".join(x) for x in rows) + "\n"
    boundary = uuid.uuid4().hex
    fields = {"request": "xmatch", "distMaxArcsec": "20", "RESPONSEFORMAT": "csv", "colRA1": "ra",
              "colDec1": "dec", "cat2": "vizier:I/355/gaiadr3", "cols2": "Source,Plx,e_Plx,Gmag,PM"}
    body = b""
    for k, v in fields.items():
        body += (f"--{boundary}\r\nContent-Disposition: form-data; name=\"{k}\"\r\n\r\n{v}\r\n").encode()
    body += (f"--{boundary}\r\nContent-Disposition: form-data; name=\"cat1\"; filename=\"cat1.csv\"\r\n"
             f"Content-Type: text/csv\r\n\r\n").encode() + buf.encode() + f"\r\n--{boundary}--\r\n".encode()
    print("  CDS XMatch -> harp2016_x_gaiadr3.csv"); sys.stdout.flush()
    b = _get(XMATCH, data=body, headers={"Content-Type": f"multipart/form-data; boundary={boundary}"})
    if not b.startswith(b"angDist"):
        raise RuntimeError("xmatch error: " + b[:400].decode(errors="replace"))
    with open(out, "wb") as f:
        f.write(b)


def fetch_gaia_lum_relation():
    """Server-side binned M_G -> log10(L/Lsun) relation (Gaia DR3 FLAME luminosities, stars within 100 pc)."""
    out = os.path.join(RAW, "gaia_mg_loglum_relation.csv")
    if cached(out):
        return
    q = ("SELECT FLOOR((s.phot_g_mean_mag+5*LOG10(s.parallax)-10)*2)/2 AS mg_bin, COUNT(*) AS n, "
         "AVG(LOG10(a.lum_flame)) AS mean_loglum, STDDEV(LOG10(a.lum_flame)) AS sd_loglum "
         "FROM gaiadr3.gaia_source AS s JOIN gaiadr3.astrophysical_parameters AS a ON s.source_id=a.source_id "
         "WHERE s.parallax>10 AND s.parallax_over_error>20 AND a.lum_flame IS NOT NULL AND s.ruwe<1.4 "
         "GROUP BY mg_bin ORDER BY mg_bin")
    params = urllib.parse.urlencode({"REQUEST": "doQuery", "LANG": "ADQL", "FORMAT": "csv", "QUERY": q})
    print("  Gaia TAP -> gaia_mg_loglum_relation.csv"); sys.stdout.flush()
    b = _get("https://gea.esac.esa.int/tap-server/tap/sync?" + params)
    if not b.startswith(b"mg_bin"):
        raise RuntimeError("Gaia TAP error: " + b[:400].decode(errors="replace"))
    with open(out, "wb") as f:
        f.write(b)


def verify_arxiv(offline=False):
    """Resolve every arXiv identifier cited in data/surveys_compilation.csv through the arXiv API and write
    data/arxiv_verification.csv (first author, year, title and journal reference as recorded by arXiv), so
    that the citations in the hand-typed compilation can be checked mechanically."""
    import json
    comp = os.path.join(DATA, "surveys_compilation.csv")
    cache = os.path.join(RAW, "arxiv_verify.json")
    rows = list(csv.DictReader(open(comp, encoding="utf8")))
    ids = sorted({m.group(1) for r in rows for m in [re.search(r"arxiv.org/abs/([\d.]+)", r["url"])] if m})
    meta = json.load(open(cache)) if os.path.exists(cache) else {}
    todo = [i for i in ids if i not in meta]
    if todo and not offline:
        for k in range(0, len(todo), 20):
            chunk = todo[k:k + 20]
            print(f"  arXiv API -> {len(chunk)} records"); sys.stdout.flush()
            try:
                xml = _get("https://export.arxiv.org/api/query?" + urllib.parse.urlencode(
                    {"id_list": ",".join(chunk), "max_results": 50}), timeout=90, tries=5, pause=20).decode()
            except RuntimeError as err:       # the arXiv API rate-limits (HTTP 429); this check is ancillary
                print("  WARNING: arXiv API unavailable, citation check incomplete; rerun later:", str(err)[-60:])
                break
            for e in re.findall(r"<entry>(.*?)</entry>", xml, flags=re.S):
                g = lambda t: re.sub(r"\s+", " ", html.unescape(re.search(r"<%s[^>]*>(.*?)</%s>" % (t, t), e, flags=re.S).group(1))).strip()
                au = re.findall(r"<name>(.*?)</name>", e)
                jr = re.search(r"<arxiv:journal_ref[^>]*>(.*?)</arxiv:journal_ref>", e, flags=re.S)
                meta[re.search(r"abs/([\d.]+)", g("id")).group(1)] = dict(
                    title=g("title"), first=au[0] if au else "", n_auth=len(au), year=g("published")[:4],
                    jr=html.unescape(jr.group(1)).strip() if jr else "")
            json.dump(meta, open(cache, "w"), indent=1)
            time.sleep(8)
    with open(os.path.join(DATA, "arxiv_verification.csv"), "w", newline="", encoding="utf8") as f:
        w = csv.writer(f)
        w.writerow(["compilation_id", "arxiv_id", "reference_as_compiled", "arxiv_resolves", "arxiv_first_author",
                    "arxiv_year_submitted", "arxiv_title", "arxiv_journal_ref"])
        n_ok = 0
        for r in rows:
            m = re.search(r"arxiv.org/abs/([\d.]+)", r["url"])
            if not m:
                continue
            o = meta.get(m.group(1))
            n_ok += o is not None
            w.writerow([r["id"], m.group(1), r["reference"], True if o is not None else "not checked"] +
                       ([o["first"], o["year"], o["title"], o["jr"]] if o else ["", "", "", ""]))
    print(f"  arXiv verification: {n_ok} of {len(ids)} cited arXiv identifiers resolve (the rest could not be checked)"
          if n_ok < len(ids) else f"  arXiv verification: {n_ok} of {len(ids)} cited arXiv identifiers resolve")


def main():
    print("[fetch] Technosearch"); fetch_technosearch()
    print("[fetch] Gaia"); fetch_gaia_lum_relation()
    print("[fetch] VizieR"); fetch_vizier()
    print("[fetch] XMatch"); xmatch_harp()
    print("[fetch] done; raw dir:", RAW)


if __name__ == "__main__":
    main()