import requests
from datetime import datetime, timedelta, timezone
import math
import os

# ── Mission window — fixed dates covering launch+1h to splashdown ──────────
# Using fixed dates means every fetch covers the entire mission arc regardless
# of when the script runs. The past/future split is done in csv_to_json.py
# based on the current time, so the display always shows what's happened and
# what's ahead.

# Launch: 2026-Apr-01 22:24 UTC
# Trajectory available in Horizons ~1h after launch
MISSION_START = datetime(2026, 4, 2, 2, 0, 0, tzinfo=timezone.utc)

# Splashdown estimate: 10 days after launch
MISSION_END   = datetime(2026, 4, 10, 22, 54, 0, tzinfo=timezone.utc)

# Step sizes
# 10-min steps over ~239 hours = ~1,434 vector points — full mission arc
VECTOR_STEP_MINUTES   = 10   # position/velocity
ELEMENTS_STEP_MINUTES = 30   # orbital elements change slowly, coarser is fine

EARTH_RADIUS_KM = 6371.0


def build_url(params):
    """
    Build the Horizons API URL as a plain concatenated string.
    
    Using requests' params= dict causes it to URL-encode the values,
    which double-encodes the single quotes Horizons requires around
    each parameter value, causing the request to fail.
    Building the query string manually keeps the quotes literal,
    exactly as the Horizons web interface sends them.
    """
    base = "https://ssd.jpl.nasa.gov/api/horizons.api"
    qs = "&".join(f"{k}={v}" for k, v in params.items())
    return f"{base}?{qs}"


def fmt_dt(dt):
    """Format datetime for Horizons — month abbreviation avoids numeric parsing issues."""
    return dt.strftime("%Y-%b-%d %H:%M")


def fetch_artemis2_vectors(start, stop, step_minutes=30):
    """Fetch state vectors (X/Y/Z position and velocity) from JPL Horizons."""

    url = build_url({
        "format":     "json",
        "COMMAND":    "'-1024'",
        "EPHEM_TYPE": "'VECTORS'",
        "CENTER":     "'500@399'",
        "START_TIME": f"'{fmt_dt(start)}'",
        "STOP_TIME":  f"'{fmt_dt(stop)}'",
        "STEP_SIZE":  f"'{step_minutes} m'",
        "VEC_TABLE":  "'2'",
        "OUT_UNITS":  "'KM-S'",
        "CSV_FORMAT": "'YES'",
        "OBJ_DATA":   "'NO'",
    })

    print(f"  Vector URL (truncated): {url[:100]}...")

    try:
        response = requests.get(url, timeout=30)
        data = response.json()
    except Exception as e:
        print(f"  Request failed: {e}")
        return None

    if "error" in data:
        print(f"  Vectors error: {data['error']}")
        return None

    records = []
    in_data = False
    for line in data["result"].split("\n"):
        if "$$SOE" in line:
            in_data = True
            continue
        if "$$EOE" in line:
            break
        if in_data and line.strip():
            parts = line.strip().split(",")
            if len(parts) >= 8:
                try:
                    x  = float(parts[2])
                    y  = float(parts[3])
                    z  = float(parts[4])
                    vx = float(parts[5])
                    vy = float(parts[6])
                    vz = float(parts[7])
                    records.append({
                        "time":      parts[1].strip(),
                        "x_km":      x,
                        "y_km":      y,
                        "z_km":      z,
                        "vx_kms":    vx,
                        "vy_kms":    vy,
                        "vz_kms":    vz,
                        "range_km":  (x**2 + y**2 + z**2)**0.5,
                        "speed_kms": (vx**2 + vy**2 + vz**2)**0.5,
                    })
                except ValueError:
                    continue

    return records


def fetch_artemis2_elements(start, stop, step_minutes=60):
    """
    Fetch osculating orbital elements.
    
    Key columns returned (CSV_FORMAT=YES):
      2=EC  eccentricity
      3=QR  periapsis distance from Earth centre (km)
      4=IN  inclination (deg)
     11=A   semi-major axis (km)
     12=AD  apoapsis distance from Earth centre (km)
    
    Note: when EC >= 1.0 the orbit is hyperbolic (trans-lunar injection or
    lunar flyby) — apogee is undefined and semi-major axis is negative.
    """
    url = build_url({
        "format":     "json",
        "COMMAND":    "'-1024'",
        "EPHEM_TYPE": "'ELEMENTS'",
        "CENTER":     "'500@399'",
        "START_TIME": f"'{fmt_dt(start)}'",
        "STOP_TIME":  f"'{fmt_dt(stop)}'",
        "STEP_SIZE":  f"'{step_minutes} m'",
        "OUT_UNITS":  "'KM-S'",
        "CSV_FORMAT": "'YES'",
        "OBJ_DATA":   "'NO'",
    })

    try:
        response = requests.get(url, timeout=30)
        data = response.json()
    except Exception as e:
        print(f"  Request failed: {e}")
        return None

    if "error" in data:
        print(f"  Elements error: {data['error']}")
        return None

    records = []
    in_data = False
    for line in data["result"].split("\n"):
        if "$$SOE" in line:
            in_data = True
            continue
        if "$$EOE" in line:
            break
        if in_data and line.strip():
            parts = line.strip().split(",")
            if len(parts) >= 13:
                try:
                    ec = float(parts[2])
                    qr = float(parts[3])
                    a  = float(parts[11])
                    ad = float(parts[12])
                    records.append({
                        "time":               parts[1].strip(),
                        "eccentricity":       ec,
                        "semi_major_axis_km": a,
                        "perigee_dist_km":    qr,
                        "apogee_dist_km":     ad,
                        "perigee_alt_km":     qr - EARTH_RADIUS_KM,
                        # apogee meaningless for hyperbolic trajectory
                        "apogee_alt_km":      (ad - EARTH_RADIUS_KM) if ec < 1.0 else None,
                        "inclination_deg":    float(parts[4]),
                        "hyperbolic":         ec >= 1.0,
                    })
                except ValueError:
                    continue

    return records


def save_to_csv(vectors, elements, filename=None):
    if filename is None:
        ts = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
        filename = f"artemis2_{ts}.csv"

    # Build element lookup keyed by time string prefix for fast nearest-match
    elem_list = elements or []

    def find_nearest_elements(t):
        """Find element record whose time is closest to t (string match first, then fallback)."""
        if not elem_list:
            return None
        t_prefix = t[:16]   # "YYYY-Mon-DD HH:MM"
        for e in elem_list:
            if e["time"][:16] == t_prefix:
                return e
        # Fallback: return the first record (elements change slowly)
        return elem_list[0] if elem_list else None

    with open(filename, "w") as f:
        f.write(
            "time_utc,"
            "x_km,y_km,z_km,"
            "vx_kms,vy_kms,vz_kms,"
            "range_from_earth_centre_km,"
            "altitude_km,"
            "speed_kms,"
            "perigee_alt_km,"
            "apogee_alt_km,"
            "perigee_dist_km,"
            "apogee_dist_km,"
            "eccentricity,"
            "semi_major_axis_km,"
            "inclination_deg\n"
        )

        for v in vectors:
            altitude = v["range_km"] - EARTH_RADIUS_KM
            elem = find_nearest_elements(v["time"])

            def ef(key, fmt=".3f"):
                val = elem.get(key) if elem else None
                return format(val, fmt) if val is not None else ""

            f.write(
                f"{v['time']},"
                f"{v['x_km']:.3f},{v['y_km']:.3f},{v['z_km']:.3f},"
                f"{v['vx_kms']:.6f},{v['vy_kms']:.6f},{v['vz_kms']:.6f},"
                f"{v['range_km']:.3f},"
                f"{altitude:.3f},"
                f"{v['speed_kms']:.6f},"
                f"{ef('perigee_alt_km')},"
                f"{ef('apogee_alt_km')},"
                f"{ef('perigee_dist_km')},"
                f"{ef('apogee_dist_km')},"
                f"{ef('eccentricity', '.8f')},"
                f"{ef('semi_major_axis_km')},"
                f"{ef('inclination_deg', '.6f')}\n"
            )

    return filename


if __name__ == "__main__":
    now   = datetime.now(timezone.utc)
    start = MISSION_START
    stop  = MISSION_END

    print(f"Fetching Artemis II full mission arc from JPL Horizons")
    print(f"  Window : {fmt_dt(start)} -> {fmt_dt(stop)} UTC")
    print(f"  ({(stop-start).days}d {(stop-start).seconds//3600}h total)")
    print(f"  Vectors: every {VECTOR_STEP_MINUTES} min (~{int((stop-start).total_seconds()/60/VECTOR_STEP_MINUTES)} points)")
    print(f"  Elements: every {ELEMENTS_STEP_MINUTES} min")

    print("\nFetching state vectors...")
    vectors = fetch_artemis2_vectors(start, stop, VECTOR_STEP_MINUTES)

    print("Fetching orbital elements...")
    elements = fetch_artemis2_elements(start, stop, ELEMENTS_STEP_MINUTES)

    if not vectors:
        print("\nERROR: No vector data returned.")
        print("The Horizons trajectory for -1024 may not cover this time window yet.")
        exit(1)

    print(f"\n  Vectors : {len(vectors)} records")
    print(f"  Elements: {len(elements) if elements else 0} records")

    filename = save_to_csv(vectors, elements)
    print(f"\nSaved to {filename}")

    # Print a sample
    print(f"\n{'Time':<32} {'Range km':>12} {'Alt km':>10} {'Speed km/s':>11}")
    print("-" * 70)
    step = max(1, len(vectors) // 8)   # show ~8 sample rows
    for v in vectors[::step]:
        alt = v["range_km"] - EARTH_RADIUS_KM
        print(f"{v['time']:<32} {v['range_km']:>12,.0f} {alt:>10,.0f} {v['speed_kms']:>11.3f}")