#!/usr/bin/env python3
"""Download every raw table used by the flyby study.

All downloads are cached in RAW (default fermi_raw/flyby/); delete a
file there to force a fresh download.  Only the Python standard library is used
for network access.  Sources:

  * Gaia DR3 gaia_source via the ESA Gaia archive TAP service
  * Gaia Catalogue of Nearby Stars (Gaia Collaboration, Smart et al. 2021),
    VizieR J/A+A/649/A6
  * The 10 parsec sample (Reyle et al. 2021), VizieR J/A+A/650/A201
  * Published encounter lists for comparison: Bailer-Jones et al. 2018
    (J/A+A/616/A37) and Bailer-Jones 2022 (J/ApJ/935/L9)
"""
import os
import sys
import time
import urllib.error
import urllib.parse
import urllib.request

RAW = os.environ.get("FLYBY_RAW", "fermi_raw/flyby")
VIZIER = "https://tapvizier.cds.unistra.fr/TAPVizieR/tap"
GAIA = "https://gea.esac.esa.int/tap-server/tap"


def _post(url, params, timeout=600):
    data = urllib.parse.urlencode(params).encode()
    req = urllib.request.Request(url, data=data)
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            return r.read(), r.geturl()
    except urllib.error.HTTPError as e:  # surface the TAP error message
        body = e.read()[:600].decode(errors="replace")
        raise RuntimeError(f"HTTP {e.code}: {body}") from None


def _get(url, timeout=600):
    with urllib.request.urlopen(url, timeout=timeout) as r:
        return r.read()


def tap_sync(base, query, out, tries=4):
    """Synchronous TAP query -> CSV file (cached)."""
    path = os.path.join(RAW, out)
    if os.path.exists(path) and os.path.getsize(path) > 0:
        print(f"[cached] {out}")
        return path
    last = None
    for k in range(tries):
        try:
            body, _ = _post(base + "/sync", dict(REQUEST="doQuery", LANG="ADQL",
                                                 FORMAT="csv", QUERY=query))
            if body.lstrip().startswith(b"<?xml") or body.lstrip().startswith(b"<VOTABLE"):
                raise RuntimeError("TAP error: " + body[:600].decode(errors="replace"))
            with open(path, "wb") as f:
                f.write(body)
            print(f"[done] {out}: {body.count(b'\n') - 1} rows")
            return path
        except Exception as e:  # network hiccups: retry
            last = e
            print(f"  retry {k + 1} for {out}: {e}")
            time.sleep(5 * (k + 1))
    raise RuntimeError(f"failed to download {out}: {last}")


def tap_async(base, query, out, poll=10, max_wait=3600):
    """Asynchronous TAP (UWS) query -> CSV file (cached)."""
    path = os.path.join(RAW, out)
    if os.path.exists(path) and os.path.getsize(path) > 0:
        print(f"[cached] {out}")
        return path
    _, job_url = _post(base + "/async", dict(REQUEST="doQuery", LANG="ADQL",
                                             FORMAT="csv", PHASE="RUN", QUERY=query))
    job_url = job_url.rstrip("/")
    print(f"  async job {job_url}")
    t0 = time.time()
    while True:
        phase = _get(job_url + "/phase").decode().strip()
        if phase == "COMPLETED":
            break
        if phase in ("ERROR", "ABORTED"):
            err = _get(job_url + "/error").decode(errors="replace")[:800]
            raise RuntimeError(f"async job {phase}: {err}")
        if time.time() - t0 > max_wait:
            raise RuntimeError("async job timed out")
        time.sleep(poll)
    body = _get(job_url + "/results/result", timeout=1800)
    with open(path, "wb") as f:
        f.write(body)
    print(f"[done] {out}: {body.count(b'\n') - 1} rows in {time.time() - t0:.0f} s")
    return path


def chunked_sync(base, template, edges, out):
    """Run `template` (with {lo},{hi} parallax bounds) in chunks and
    concatenate the CSV results into one cached file."""
    path = os.path.join(RAW, out)
    if os.path.exists(path) and os.path.getsize(path) > 0:
        print(f"[cached] {out}")
        return path
    parts = []
    for lo, hi in zip(edges[:-1], edges[1:]):
        part = out.replace(".csv", f".part_{lo}_{hi}.csv")
        tap_sync(base, template.format(lo=lo, hi=hi), part)
        parts.append(os.path.join(RAW, part))
    n = 0
    with open(path, "w") as f:
        for i, pth in enumerate(parts):
            with open(pth) as g:
                header = g.readline()
                if i == 0:
                    f.write(header)
                for line in g:
                    f.write(line)
                    n += 1
    for pth in parts:
        os.remove(pth)
    print(f"[done] {out}: {n} rows total")
    return path


# ----------------------------------------------------------------------------
# Queries
# ----------------------------------------------------------------------------
# (1) Every Gaia DR3 source with a radial velocity and parallax > 8 mas.
#     Used for the local velocity distribution and the RV-completeness map.
Q_RV125 = """
SELECT source_id, ra, dec, l, b, parallax, parallax_error, pmra, pmdec,
       radial_velocity, radial_velocity_error, phot_g_mean_mag, bp_rp, grvs_mag,
       ruwe, rv_expected_sig_to_noise, rv_nb_transits, rv_template_teff,
       rv_renormalised_gof, rv_amplitude_robust
FROM gaiadr3.gaia_source
WHERE parallax > {lo} AND parallax <= {hi} AND radial_velocity IS NOT NULL
"""

# (2) Encounter candidates over the whole RV catalogue.  Linear-motion
#     perihelion distance  d_ph = r * vT / v  with r = 1000/parallax [pc],
#     vT = 4.740470 * mu / parallax [km/s].  Keep d_ph < 15 pc (no time cut;
#     the time window is applied in the analysis).  Restricted to parallax
#     > 2 mas (500 pc): a star farther away would need v > 250 km/s to have
#     its perihelion inside the +-2 Myr window where linear motion is used.
Q_ENC = """
SELECT source_id, ra, dec, l, b, parallax, parallax_error, pmra, pmra_error,
       pmdec, pmdec_error, radial_velocity, radial_velocity_error,
       phot_g_mean_mag, bp_rp, grvs_mag, ruwe, rv_expected_sig_to_noise,
       rv_nb_transits, rv_template_teff, rv_renormalised_gof,
       rv_amplitude_robust, ipd_frac_multi_peak, astrometric_params_solved,
       non_single_star
FROM gaiadr3.gaia_source
WHERE radial_velocity IS NOT NULL
  AND parallax > {lo} AND parallax <= {hi} AND parallax_over_error > 5
  AND (1000.0/parallax) * (4.740470*SQRT(pmra*pmra+pmdec*pmdec)/parallax)
      / SQRT(POWER(4.740470*SQRT(pmra*pmra+pmdec*pmdec)/parallax, 2)
             + radial_velocity*radial_velocity) < 15
"""

# (3) GCNS main table (331 312 sources within 100 pc).
Q_GCNS = """
SELECT GaiaEDR3, RA_ICRS, DE_ICRS, Plx, e_Plx, pmRA, pmDE, "Gmag", BPmag, RPmag,
       RUWE, RV, e_RV, GCNSprob, WDprob, Dist16, Dist50, Dist84,
       xcoord50, ycoord50, zcoord50
FROM "J/A+A/649/A6/table1c"
"""
Q_GCNS_BIN = 'SELECT * FROM "J/A+A/649/A6/table3"'
Q_10PC = 'SELECT * FROM "J/A+A/650/A201/tablea1"'
Q_BJ18 = 'SELECT * FROM "J/A+A/616/A37/table23"'
Q_BJ22 = 'SELECT * FROM "J/ApJ/935/L9/table12"'


def main():
    os.makedirs(RAW, exist_ok=True)
    tap_sync(VIZIER, Q_10PC, "tenpc_reyle2021.csv")
    tap_sync(VIZIER, Q_BJ18, "bailerjones2018_table23.csv")
    tap_sync(VIZIER, Q_BJ22, "bailerjones2022_table12.csv")
    tap_sync(VIZIER, Q_GCNS_BIN, "gcns_resolved_binaries.csv")
    tap_sync(VIZIER, Q_GCNS, "gcns_table1c.csv")   # 331 312 rows, ~45 MB
    # The Gaia archive's synchronous endpoint is used in parallax chunks (the
    # anonymous async queue was observed to stall for > 15 minutes).
    chunked_sync(GAIA, Q_RV125, [8, 9, 10, 12, 15, 20, 1e6], "gaia_dr3_rv_plx8.csv")
    chunked_sync(GAIA, Q_ENC, [2, 2.5, 3, 4, 6, 1e6], "gaia_dr3_encounter_candidates.csv")


if __name__ == "__main__":
    sys.exit(main())