"""Shared helpers for the dim_stars study (paths, TAP / XMatch access, SFD lookup).

Only numpy / pandas / astropy / urllib are used.  Nothing here disables TLS
verification; the proxy and CA bundle come from the environment.
"""
import io
import os
import time
import urllib.error
import urllib.parse
import urllib.request
import uuid

import numpy as np
import pandas as pd

STUDY = "studies/dim_stars"
RAW = os.environ.get("DIM_STARS_RAW", "fermi_raw/dim_stars")
# Two sets of astrometric cuts are carried through the whole pipeline (see dimlib.CUTSETS):
#   "final"  RUWE < 1.2, parallax/error >= 50   (tightened after a first look at the faint tail)
#   "orig"   RUWE < 1.4, parallax/error >= 20   (the cuts declared before looking at the data)
# The set is chosen with the environment variable DIM_STARS_CUTSET; run_all.py runs both.
CUTSET = os.environ.get("DIM_STARS_CUTSET", "final")
assert CUTSET in ("final", "orig"), CUTSET
TAG = "" if CUTSET == "final" else "_origcuts"
DATA_MAIN = os.path.join(STUDY, "data")                                  # products of the "final" cut set
DATA_ORIG = os.path.join(DATA_MAIN, "origcuts")                          # same products, "orig" cut set
DATA = DATA_MAIN if CUTSET == "final" else DATA_ORIG
FIGS = os.path.join(STUDY, "figures")
WORK = os.path.join(RAW, "work")          # large intermediate tables (not in shared folder)
for _d in (RAW, DATA_MAIN, DATA_ORIG, FIGS, WORK):
    os.makedirs(_d, exist_ok=True)


def dpkl(survey, cutset=None):
    """Path of the per-star deficit table of a survey for a cut set (default: the active one)."""
    tag = TAG if cutset is None else ("" if cutset == "final" else "_origcuts")
    return os.path.join(WORK, "%s_D%s.pkl" % (survey, tag))


def cache(name):
    """Path of a small network cache in RAW that depends on the active cut set."""
    root, ext = os.path.splitext(name)
    return os.path.join(RAW, root + TAG + ext)

VIZIER_TAP = "https://tapvizier.cds.unistra.fr/TAPVizieR/tap/sync"
GAIA_TAP = "https://gea.esac.esa.int/tap-server/tap"
IRSA_TAP = "https://irsa.ipac.caltech.edu/TAP/sync"
XMATCH = "http://cdsxmatch.u-strasbg.fr/xmatch/api/v1/sync"
SFD_URL = "https://raw.githubusercontent.com/kbarbary/sfddata/master/SFD_dust_4096_%s.fits"


def log(*a):
    print(time.strftime("[%H:%M:%S]"), *a, flush=True)


def _post(url, fields, timeout=900, retries=4):
    data = urllib.parse.urlencode(fields).encode()
    last = None
    for k in range(retries):
        try:
            req = urllib.request.Request(url, data=data)
            with urllib.request.urlopen(req, timeout=timeout) as r:
                return r.read()
        except Exception as e:  # noqa
            last = e
            log("   request failed (%s); retry %d" % (repr(e)[:120], k + 1))
            time.sleep(5 + 10 * k)
    raise RuntimeError("request failed after retries: %r" % (last,))


def tap_sync(url, adql, maxrec=3000000, timeout=900, fmt="csv", retries=4):
    """Run a synchronous TAP query and return the raw CSV bytes."""
    out = _post(url, dict(REQUEST="doQuery", LANG="ADQL", FORMAT=fmt, MAXREC=str(maxrec), QUERY=adql),
                timeout=timeout, retries=retries)
    if out[:5] == b"<?xml":
        raise RuntimeError("TAP error: " + out[:1500].decode(errors="replace"))
    return out


def tap_sync_to_file(url, adql, path, **kw):
    if os.path.exists(path) and os.path.getsize(path) > 0:
        return path
    out = tap_sync(url, adql, **kw)
    tmp = path + ".part"
    with open(tmp, "wb") as f:
        f.write(out)
    os.replace(tmp, path)
    return path


def gaia_async_to_file(adql, path, poll=15, timeout=7200):
    """Run an asynchronous Gaia-archive TAP job and save the CSV result."""
    if os.path.exists(path) and os.path.getsize(path) > 0:
        return path
    data = urllib.parse.urlencode(dict(REQUEST="doQuery", LANG="ADQL", FORMAT="csv", PHASE="RUN",
                                       QUERY=adql)).encode()
    req = urllib.request.Request(GAIA_TAP + "/async", data=data)
    with urllib.request.urlopen(req, timeout=300) as r:
        joburl = r.geturl()
    log("   gaia async job", joburl)
    t0 = time.time()
    while True:
        with urllib.request.urlopen(joburl + "/phase", timeout=120) as r:
            phase = r.read().decode().strip()
        if phase == "COMPLETED":
            break
        if phase in ("ERROR", "ABORTED"):
            with urllib.request.urlopen(joburl + "/error", timeout=120) as r:
                raise RuntimeError("Gaia job failed: " + r.read()[:800].decode(errors="replace"))
        if time.time() - t0 > timeout:
            raise RuntimeError("Gaia job timeout")
        time.sleep(poll)
    tmp = path + ".part"
    with urllib.request.urlopen(joburl + "/results/result", timeout=3600) as r, open(tmp, "wb") as f:
        while True:
            b = r.read(1 << 20)
            if not b:
                break
            f.write(b)
    os.replace(tmp, path)
    return path


def xmatch_upload(df, ra, dec, cat2, radius, cols2=None, timeout=1800, retries=4):
    """Positional cross-match of a local table with a VizieR table via CDS XMatch.
    df must contain only the columns to be uploaded.  Returns a DataFrame."""
    boundary = uuid.uuid4().hex
    csv_bytes = df.to_csv(index=False).encode()
    fields = dict(request="xmatch", distMaxArcsec=str(radius), RESPONSEFORMAT="csv",
                  colRA1=ra, colDec1=dec, cat2=cat2, maxRec="50000000")
    if cols2:
        fields["cols2"] = cols2
    body = b""
    for k, v in fields.items():
        body += ("--%s\r\nContent-Disposition: form-data; name=\"%s\"\r\n\r\n%s\r\n" % (boundary, k, v)).encode()
    body += ("--%s\r\nContent-Disposition: form-data; name=\"cat1\"; filename=\"up.csv\"\r\n"
             "Content-Type: text/csv\r\n\r\n" % boundary).encode() + csv_bytes + b"\r\n"
    body += ("--%s--\r\n" % boundary).encode()
    last = None
    for k in range(retries):
        try:
            req = urllib.request.Request(XMATCH, data=body,
                                         headers={"Content-Type": "multipart/form-data; boundary=" + boundary})
            with urllib.request.urlopen(req, timeout=timeout) as r:
                out = r.read()
            if out[:5] == b"<?xml":
                raise RuntimeError("XMatch error: " + out[:3000].decode(errors="replace")[-1500:])
            return pd.read_csv(io.BytesIO(out), low_memory=False)
        except Exception as e:  # noqa
            last = e
            log("   xmatch failed (%s); retry %d" % (repr(e)[:160], k + 1))
            time.sleep(10 + 20 * k)
    raise RuntimeError("xmatch failed: %r" % (last,))


def download(url, path, timeout=600):
    if os.path.exists(path) and os.path.getsize(path) > 0:
        return path
    tmp = path + ".part"
    with urllib.request.urlopen(url, timeout=timeout) as r, open(tmp, "wb") as f:
        while True:
            b = r.read(1 << 20)
            if not b:
                break
            f.write(b)
    os.replace(tmp, path)
    return path


# ----------------------------------------------------------------------------
# SFD (Schlegel, Finkbeiner & Davis 1998) E(B-V) lookup, nearest pixel.
# Lambert projection as in the original SFD FITS headers (4096x4096 per pole).
# ----------------------------------------------------------------------------
_SFD = {}


def sfd_ebv(l_deg, b_deg):
    from astropy.io import fits
    l = np.radians(np.asarray(l_deg, float))
    b = np.radians(np.asarray(b_deg, float))
    out = np.full(l.shape, np.nan)
    for pole, n in (("ngp", 1), ("sgp", -1)):
        m = (b >= 0) if n == 1 else (b < 0)
        if not m.any():
            continue
        if pole not in _SFD:
            p = download(SFD_URL % pole, os.path.join(RAW, "SFD_dust_4096_%s.fits" % pole))
            _SFD[pole] = fits.getdata(p).astype(np.float32)
        img = _SFD[pole]
        x = 2048.0 * np.sqrt(1.0 - n * np.sin(b[m])) * np.cos(l[m]) + 2047.5
        y = -2048.0 * n * np.sqrt(1.0 - n * np.sin(b[m])) * np.sin(l[m]) + 2047.5
        ix = np.clip(np.rint(x).astype(int), 0, 4095)
        iy = np.clip(np.rint(y).astype(int), 0, 4095)
        out[m] = img[iy, ix]
    return out


def to_epoch2000(ra, dec, pmra, pmdec, epoch=2016.0):
    """Propagate Gaia DR3 (epoch 2016.0) positions to epoch 2000.0 (linear, small-angle)."""
    dt = 2000.0 - epoch
    pmra = np.where(np.isfinite(pmra), pmra, 0.0)
    pmdec = np.where(np.isfinite(pmdec), pmdec, 0.0)
    dec2 = dec + pmdec * dt / 3.6e6
    ra2 = ra + pmra * dt / 3.6e6 / np.cos(np.radians(dec))
    return ra2 % 360.0, dec2