"""
fetch_artemis2.py
Fetches Artemis II trajectory and orbital elements from JPL Horizons
and writes artemis2_data.json to the current directory.

Run once manually to test, then schedule with cron every 10 minutes:
  crontab -e
  */10 * * * * cd /path/to/your/webroot && python3 fetch_artemis2.py

Requires: requests
  pip3 install requests
"""

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

EARTH_RADIUS_KM = 6371.0
HORIZONS_URL    = 'https://ssd.jpl.nasa.gov/api/horizons.api'
OUTPUT_FILE     = 'artemis2_data.json'

# Horizons rejects the request if START_TIME is before the trajectory exists.
# Clamp start to the first available epoch - never request pre-launch data.
LAUNCH_TIME  = datetime(2026, 4, 2, 2, 1, 0, tzinfo=timezone.utc)   # 2 mins past first epoch
HOURS_AHEAD  = 240   # 10 days forward from now

# Step sizes
VECTOR_STEP     = '1 h'   # 1-hour position steps (~250 points over mission)
ELEMENTS_STEP   = '6 h'   # 6-hour orbital element steps


def horizons_request(params):
    """Make a request to JPL Horizons and return parsed JSON.
    
    Builds the URL as a plain string to avoid requests double-encoding
    the single-quote delimiters that Horizons requires around values.
    """
    # Build query string manually - Horizons needs literal single quotes
    # around most values, which requests.get(params=...) would double-encode
    qs = '&'.join(f"{k}={v}" for k, v in params.items())
    url = f"{HORIZONS_URL}?{qs}"
    print(f"\n  DEBUG URL: {url[:120]}...")
    try:
        resp = requests.get(url, timeout=30)
        resp.raise_for_status()
        data = resp.json()
        if 'error' in data:
            print(f"  Horizons error: {data['error']}")
            return None
        return data
    except requests.exceptions.Timeout:
        print("  Request timed out.")
        return None
    except requests.exceptions.RequestException as e:
        print(f"  Request failed: {e}")
        return None


def parse_vectors(result_text):
    """Parse CSV vector output from Horizons $$SOE...$$EOE block."""
    records = []
    in_data = False
    for line in result_text.split('\n'):
        if '$$SOE' in line:
            in_data = True
            continue
        if '$$EOE' in line:
            break
        if not in_data or not line.strip():
            continue
        parts = line.split(',')
        if len(parts) < 8:
            continue
        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':      round(x, 3),
                'y_km':      round(y, 3),
                'z_km':      round(z, 3),
                'vx_kms':    round(vx, 6),
                'vy_kms':    round(vy, 6),
                'vz_kms':    round(vz, 6),
                'range_km':  round(math.sqrt(x*x + y*y + z*z), 3),
                'speed_kms': round(math.sqrt(vx*vx + vy*vy + vz*vz), 6),
            })
        except (ValueError, IndexError):
            continue
    return records


def parse_elements(result_text):
    """Parse CSV orbital elements from Horizons $$SOE...$$EOE block.

    CSV columns (0-indexed):
      0  JDTDB
      1  Cal_Date
      2  EC  (eccentricity)
      3  QR  (periapsis distance, km from Earth centre)
      4  IN  (inclination, deg)
      5  OM  (longitude of ascending node, deg)
      6  W   (argument of periapsis, deg)
      7  Tp  (time of periapsis passage)
      8  N   (mean motion, deg/s)
      9  MA  (mean anomaly, deg)
      10 TA  (true anomaly, deg)
      11 A   (semi-major axis, km)
      12 AD  (apoapsis distance, km from Earth centre)
      13 PR  (sidereal orbit period, s)
    """
    records = []
    in_data = False
    for line in result_text.split('\n'):
        if '$$SOE' in line:
            in_data = True
            continue
        if '$$EOE' in line:
            break
        if not in_data or not line.strip():
            continue
        parts = line.split(',')
        if len(parts) < 13:
            continue
        try:
            ec = float(parts[2])
            qr = float(parts[3])
            a  = float(parts[11])
            ad = float(parts[12])
            records.append({
                'time':               parts[1].strip(),
                'eccentricity':       round(ec, 8),
                'semi_major_axis_km': round(a, 3),
                'perigee_dist_km':    round(qr, 3),
                'apogee_dist_km':     round(ad, 3),
                'perigee_alt_km':     round(qr - EARTH_RADIUS_KM, 3),
                # apogee is meaningless when ec >= 1 (hyperbolic trajectory)
                'apogee_alt_km':      round(ad - EARTH_RADIUS_KM, 3) if ec < 1.0 else None,
                'inclination_deg':    round(float(parts[4]), 6),
                'hyperbolic':         ec >= 1.0,
            })
        except (ValueError, IndexError):
            continue
    return records


def fmt_dt(dt):
    # Horizons strongly prefers month abbreviation format e.g. 2026-Apr-02 18:00
    # Numeric month format (2026-04-02) can cause rejection on some queries
    return dt.strftime('%Y-%b-%d %H:%M')


def run():
    now   = datetime.now(timezone.utc)
    # Start from launch time (trajectory start) or 48h ago, whichever is later
    start = max(LAUNCH_TIME, now - timedelta(hours=48))
    stop  = now + timedelta(hours=HOURS_AHEAD)

    print(f"Fetching Artemis II data from JPL Horizons...")
    print(f"  Window: {fmt_dt(start)} -> {fmt_dt(stop)} UTC")

    # ── Vectors ──
    print("  Requesting state vectors...", end=' ', flush=True)
    v_data = horizons_request({
        '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"'{VECTOR_STEP}'",
        'VEC_TABLE':  "'2'",
        'OUT_UNITS':  "'KM-S'",
        'CSV_FORMAT': "'YES'",
        'OBJ_DATA':   "'NO'",
    })

    if not v_data:
        print("FAILED")
        return False

    vectors = parse_vectors(v_data['result'])
    if not vectors:
        print("FAILED - no data in response (trajectory not yet available)")
        return False
    print(f"OK ({len(vectors)} points)")

    # ── Orbital elements ──
    print("  Requesting orbital elements...", end=' ', flush=True)
    e_data = horizons_request({
        '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"'{ELEMENTS_STEP}'",
        'OUT_UNITS':  "'KM-S'",
        'CSV_FORMAT': "'YES'",
        'OBJ_DATA':   "'NO'",
    })

    elements = []
    if e_data:
        elements = parse_elements(e_data['result'])
        print(f"OK ({len(elements)} points)")
    else:
        print("FAILED (continuing without elements)")

    # ── Find current position index ──
    now_ms = now.timestamp() * 1000
    current_idx = 0
    best_diff = float('inf')
    for i, v in enumerate(vectors):
        # Horizons time format: "A.D. 2026-Apr-01 18:00:00.0000 TDB"
        # Strip prefix and suffix for parsing
        t_str = v['time']
        try:
            # Remove "A.D. " prefix and " TDB" / " UT" suffix
            clean = t_str.replace('A.D. ', '').split(' TDB')[0].split(' UT')[0].strip()
            t = datetime.strptime(clean, '%Y-%b-%d %H:%M:%S.%f').replace(tzinfo=timezone.utc)
            diff = abs(t.timestamp() - now.timestamp())
            if diff < best_diff:
                best_diff = diff
                current_idx = i
        except Exception:
            continue

    current = vectors[current_idx] if vectors else None

    # ── Write JSON ──
    output = {
        'fetched_at':   now.isoformat(),
        'current_idx':  current_idx,
        'current':      current,
        'vectors':      vectors,
        'elements':     elements,
        'mission_info': {
            'name':        'Artemis II',
            'spacecraft':  'Orion / Integrity',
            'horizons_id': '-1024',
            'launched':    '2026-04-01T22:24:00Z',
        }
    }

    tmp = OUTPUT_FILE + '.tmp'
    with open(tmp, 'w') as f:
        json.dump(output, f, separators=(',', ':'))
    os.replace(tmp, OUTPUT_FILE)  # atomic replace - safe for live webservers

    alt = current['range_km'] - EARTH_RADIUS_KM if current else 0
    print(f"  Written {OUTPUT_FILE}")
    print(f"  Current position: range={current['range_km']:,.0f} km, altitude={alt:,.0f} km, speed={current['speed_kms']:.3f} km/s")
    return True


if __name__ == '__main__':
    success = run()
    if not success:
        print("\nFetch failed. If this is shortly after launch, the trajectory")
        print("data (-1024) may not yet be published in JPL Horizons.")
        print("Try again in a few hours.")
