#!/usr/bin/env python3
"""Fetch and store TRAPPIST-1 data from the NASA Exoplanet Archive.

This script is intentionally narrow:
1. It fetches all `pscomppars` rows for the `TRAPPIST-1` host star.
2. It stores the raw NASA response unchanged for provenance.
3. It writes a normalized server-side payload containing one star record and
   one planet record per returned row.

The script uses `pscomppars` for both the star and planet data. Host-star
fields are repeated on each planet row, so a separate `stellarhosts` query is
not required for the initial implementation.
"""

from __future__ import annotations

import argparse
import json
import math
import os
import sys
import urllib.parse
import urllib.request
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

ASTROPY_CACHE_ROOT = Path(__file__).resolve().parent / ".astropy-cache"
ASTROPY_CACHE_ROOT.mkdir(parents=True, exist_ok=True)
os.environ.setdefault("ASTROPY_CACHE_DIR", str(ASTROPY_CACHE_ROOT / "cache"))

try:
    from astropy.time import Time
except ImportError:  # pragma: no cover - exercised only when dependency is missing at runtime.
    Time = None


NASA_TAP_SYNC_URL = "https://exoplanetarchive.ipac.caltech.edu/TAP/sync"
RUNTIME_ENDPOINTS_CONFIG = Path(__file__).resolve().parents[1] / "config" / "runtime_endpoints.json"
EARTH_GRAVITY_MPS2 = 9.80665
EARTH_RADIUS_M = 6_371_000.0
EARTH_MASS_KG = 5.9722e24
AU_M = 149_597_870_700.0
SOLAR_RADIUS_M = 695_700_000.0
SOLAR_MASS_KG = 1.98847e30
SOLAR_EFFECTIVE_TEMPERATURE_K = 5_772.0
UTC = timezone.utc


@dataclass(frozen=True)
class FetchConfig:
    """Configuration for the archive fetch."""

    hostname: str
    output_dir: Path
    timeout_seconds: float
    tap_sync_url: str


def parse_args() -> FetchConfig:
    """Parse command-line arguments into a typed configuration."""

    parser = argparse.ArgumentParser(
        description=(
            "Fetch TRAPPIST-1 planetary-system composite parameters from the "
            "NASA Exoplanet Archive and store raw plus canonical JSON outputs."
        )
    )
    parser.add_argument(
        "--hostname",
        default="TRAPPIST-1",
        help="Host star name to query from the pscomppars table.",
    )
    parser.add_argument(
        "--output-dir",
        default=str(Path(__file__).resolve().parent),
        help="Directory where JSON outputs will be written.",
    )
    parser.add_argument(
        "--timeout-seconds",
        type=float,
        default=30.0,
        help="HTTP timeout for the NASA archive request.",
    )
    parser.add_argument(
        "--tap-sync-url",
        default=None,
        help=(
            "NASA TAP sync URL. Defaults to TR1_NASA_TAP_SYNC_URL, then "
            "config/runtime_endpoints.json, then the official NASA endpoint."
        ),
    )
    args = parser.parse_args()
    return FetchConfig(
        hostname=args.hostname,
        output_dir=Path(args.output_dir),
        timeout_seconds=args.timeout_seconds,
        tap_sync_url=resolve_tap_sync_url(args.tap_sync_url),
    )


def load_runtime_endpoints_config() -> dict[str, Any]:
    """Load environment-specific endpoint config when present."""

    if not RUNTIME_ENDPOINTS_CONFIG.exists():
        return {}
    payload = json.loads(RUNTIME_ENDPOINTS_CONFIG.read_text(encoding="utf-8"))
    if not isinstance(payload, dict):
        raise ValueError("runtime_endpoints.json must contain a JSON object.")
    return payload


def resolve_tap_sync_url(cli_value: str | None) -> str:
    """Resolve the NASA TAP sync URL from CLI, env, config, or default."""

    if cli_value:
        return cli_value

    env_value = os.environ.get("TR1_NASA_TAP_SYNC_URL")
    if env_value:
        return env_value

    config = load_runtime_endpoints_config()
    nasa_archive_config = config.get("nasa_archive", {})
    if isinstance(nasa_archive_config, dict):
        configured_url = nasa_archive_config.get("tap_sync_url")
        if isinstance(configured_url, str) and configured_url:
            return configured_url

    return NASA_TAP_SYNC_URL


def build_pscomppars_query(hostname: str) -> str:
    """Build the SQL query for the `pscomppars` table."""

    escaped_hostname = hostname.replace("'", "''")
    return (
        "select * "
        "from pscomppars "
        f"where hostname = '{escaped_hostname}' "
        "order by pl_letter"
    )


def build_tap_url(tap_sync_url: str, query: str) -> str:
    """Build the NASA TAP URL for a JSON response."""

    encoded_query = urllib.parse.urlencode({"query": query, "format": "json"})
    return f"{tap_sync_url}?{encoded_query}"


def fetch_rows(url: str, timeout_seconds: float) -> list[dict[str, Any]]:
    """Fetch JSON rows from the NASA Exoplanet Archive TAP endpoint."""

    request = urllib.request.Request(
        url,
        headers={
            "User-Agent": "Trappist1 NASA Archive Fetcher/1.0",
            "Accept": "application/json",
        },
    )
    with urllib.request.urlopen(request, timeout=timeout_seconds) as response:
        payload = response.read().decode("utf-8")
    rows = json.loads(payload)
    if not isinstance(rows, list):
        raise ValueError("NASA archive returned a non-list JSON payload.")
    for row in rows:
        if not isinstance(row, dict):
            raise ValueError("NASA archive returned a row that is not an object.")
    return rows


def require_value(row: dict[str, Any], key: str) -> Any:
    """Return a required value from a row or raise a descriptive error."""

    value = row.get(key)
    if value is None:
        raise ValueError(f"Required field '{key}' is missing from NASA row.")
    return value


def to_float(value: Any) -> float | None:
    """Convert a NASA field to float when present."""

    if value is None:
        return None
    return float(value)


def to_string(value: Any) -> str | None:
    """Convert a NASA field to string when present."""

    if value is None:
        return None
    return str(value)


def normalize_angle_degrees(angle_degrees: float | None) -> float | None:
    """Normalize an angle to the [0, 360) range when present."""

    if angle_degrees is None:
        return None
    normalized = math.fmod(angle_degrees, 360.0)
    if normalized < 0.0:
        normalized += 360.0
    return normalized


def derive_surface_gravity_earth(mass_earth: float | None, radius_earth: float | None) -> float | None:
    """Calculate surface gravity relative to Earth."""

    if mass_earth is None or radius_earth is None or radius_earth == 0.0:
        return None
    return mass_earth / (radius_earth * radius_earth)


def derive_orbital_degrees_per_second(period_seconds: float | None) -> float | None:
    """Calculate orbital angular velocity in degrees per second."""

    if period_seconds is None or period_seconds == 0.0:
        return None
    return 360.0 / period_seconds


def derive_rotation_axis_tilt_from_sky_inclination(inclination_deg: float | None) -> float | None:
    """Convert NASA sky inclination into the simple tilt convention used in Unity.

    The existing Unity prototype uses `90 - inclination` as a practical tilt
    value for the orbit wrapper rather than attempting a full orbital-frame
    reconstruction.
    """

    if inclination_deg is None:
        return None
    return 90.0 - inclination_deg


def derive_relative_ratio(value: float | None, baseline: float) -> float | None:
    """Calculate a unitless ratio relative to a baseline."""

    if value is None or baseline == 0.0:
        return None
    return value / baseline


def calculate_most_recent_mid_transit_utc(
    reference_mid_transit_bjd_tdb: float | None,
    orbital_period_days: float | None,
    now_utc: datetime,
) -> tuple[str | None, float | None]:
    """Calculate the most recent mid-transit timestamp in UTC and Unix seconds.

    Scientific accuracy matters for this project, so this function deliberately
    avoids the earlier "treat BJD-TDB as a plain Julian day" approximation.

    The NASA archive provides `pl_tranmid` as a barycentric Julian date in the
    TDB time scale. To preserve that meaning:

    1. Interpret the archive value as `JD(TDB)`.
    2. Convert the current UTC refresh moment into the same `TDB` scale.
    3. Step forward by whole orbital periods while remaining in `TDB`.
    4. Convert the resulting timestamp from `TDB` to `UTC`.
    5. Export both ISO-8601 UTC text and Unix seconds for the client.

    This is the minimum honest pipeline if the output field is labeled as UTC.
    """

    if reference_mid_transit_bjd_tdb is None or orbital_period_days is None or orbital_period_days <= 0.0:
        return None, None

    if Time is None:
        raise RuntimeError(
            "Scientific transit-time conversion requires astropy. "
            "Install it in the runtime environment before generating UTC/Unix transit anchors."
        )

    # `pl_tranmid` from `pscomppars` is documented as a barycentric Julian date
    # in TDB. We preserve that meaning explicitly instead of collapsing it into
    # Unix-epoch arithmetic too early.
    reference_tdb = Time(reference_mid_transit_bjd_tdb, format="jd", scale="tdb")

    # Convert the current refresh moment to an astropy UTC time, then into TDB
    # so that the elapsed-period arithmetic is performed inside one consistent
    # time scale.
    now_time_utc = Time(now_utc, format="datetime", scale="utc")
    now_time_tdb = now_time_utc.tdb

    elapsed_periods = math.floor((now_time_tdb.jd - reference_tdb.jd) / orbital_period_days)
    most_recent_tdb = Time(
        reference_tdb.jd + (elapsed_periods * orbital_period_days),
        format="jd",
        scale="tdb",
    )

    # Only after the TDB stepping is complete do we convert the result to UTC
    # and then to Unix seconds for the Unity client.
    most_recent_utc = most_recent_tdb.utc
    return most_recent_utc.isot + "Z", float(most_recent_utc.unix)


def extract_star_record(first_row: dict[str, Any], hostname: str) -> dict[str, Any]:
    """Build a normalized star record from the shared host-star fields."""

    stellar_radius_solar = to_float(first_row.get("st_rad"))
    stellar_mass_solar = to_float(first_row.get("st_mass"))
    effective_temperature_k = to_float(first_row.get("st_teff"))
    luminosity_log10_solar = to_float(first_row.get("st_lum"))
    radius_m = stellar_radius_solar * SOLAR_RADIUS_M if stellar_radius_solar is not None else None
    mass_kg = stellar_mass_solar * SOLAR_MASS_KG if stellar_mass_solar is not None else None

    return {
        "id": hostname.lower().replace(" ", "-"),
        "name": require_value(first_row, "hostname"),
        "kind": "star",
        "physical": {
            "mass_kg": mass_kg,
            "radius_m": radius_m,
            "effective_temperature_k": effective_temperature_k,
            "luminosity_log10_solar": luminosity_log10_solar,
        },
    }


def extract_planet_record(row: dict[str, Any], star_id: str, now_utc: datetime) -> dict[str, Any]:
    """Build a normalized planet record from one `pscomppars` row."""

    radius_earth = to_float(row.get("pl_rade"))
    best_mass_earth = to_float(row.get("pl_bmasse"))
    surface_gravity_earth = derive_surface_gravity_earth(best_mass_earth, radius_earth)
    orbital_period_days = to_float(row.get("pl_orbper"))
    orbital_period_seconds = orbital_period_days * 86400.0 if orbital_period_days is not None else None
    semi_major_axis_au = to_float(row.get("pl_orbsmax"))
    semi_major_axis_m = semi_major_axis_au * AU_M if semi_major_axis_au is not None else None
    transit_midpoint_bjd_tdb = to_float(row.get("pl_tranmid"))
    inclination_deg = to_float(row.get("pl_orbincl"))
    argument_of_periastron_deg = normalize_angle_degrees(to_float(row.get("pl_orblper")))
    radius_m = radius_earth * EARTH_RADIUS_M if radius_earth is not None else None
    mass_kg = best_mass_earth * EARTH_MASS_KG if best_mass_earth is not None else None
    surface_gravity_mps2 = (
        surface_gravity_earth * EARTH_GRAVITY_MPS2
        if surface_gravity_earth is not None
        else None
    )
    last_mid_transit_utc, last_mid_transit_unix_seconds = calculate_most_recent_mid_transit_utc(
        transit_midpoint_bjd_tdb,
        orbital_period_days,
        now_utc,
    )

    planet_name = require_value(row, "pl_name")
    planet_letter = require_value(row, "pl_letter")

    return {
        "id": planet_name.lower().replace(" ", "-"),
        "name": planet_name,
        "letter": planet_letter,
        "kind": "planet",
        "parent_id": star_id,
        "physical": {
            "radius_m": radius_m,
            "mass_kg": mass_kg,
            "equilibrium_temperature_k": to_float(row.get("pl_eqt")),
            "surface_gravity_mps2": surface_gravity_mps2,
        },
        "orbit": {
            "semi_major_axis_m": semi_major_axis_m,
            "orbital_period_seconds": orbital_period_seconds,
            "orbital_degrees_per_second": derive_orbital_degrees_per_second(orbital_period_seconds),
            "eccentricity": to_float(row.get("pl_orbeccen")),
            "inclination_deg": inclination_deg,
            "unity_orbit_tilt_deg": derive_rotation_axis_tilt_from_sky_inclination(inclination_deg),
            "argument_of_periastron_deg": argument_of_periastron_deg,
            "last_mid_transit_unix_seconds": last_mid_transit_unix_seconds,
        },
    }


def build_canonical_payload(rows: list[dict[str, Any]], hostname: str, now_utc: datetime) -> dict[str, Any]:
    """Construct the normalized star-plus-planets payload."""

    if not rows:
        raise ValueError(f"No pscomppars rows were returned for hostname '{hostname}'.")

    first_row = rows[0]
    star = extract_star_record(first_row, hostname)
    planets = [extract_planet_record(row, star["id"], now_utc) for row in rows]

    return {
        "system_name": to_string(first_row.get("sy_name")) or hostname,
        "star": star,
        "planets": planets,
    }


def ensure_consistent_star_fields(rows: list[dict[str, Any]]) -> list[str]:
    """Report any shared star fields that differ across the returned rows."""

    fields_to_check = [
        "hostname",
        "sy_name",
        "st_teff",
        "st_rad",
        "st_mass",
        "st_lum",
        "st_age",
        "sy_dist",
    ]
    mismatches: list[str] = []
    if not rows:
        return mismatches

    first_row = rows[0]
    for field in fields_to_check:
        first_value = first_row.get(field)
        for row in rows[1:]:
            if row.get(field) != first_value:
                mismatches.append(field)
                break
    return mismatches


def write_json(path: Path, payload: Any) -> None:
    """Write a JSON file with deterministic formatting."""

    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")


def main() -> int:
    """Run the fetch/store workflow."""

    config = parse_args()
    query = build_pscomppars_query(config.hostname)
    url = build_tap_url(config.tap_sync_url, query)
    fetch_moment_utc = datetime.now(UTC).replace(microsecond=0)
    fetched_at_utc = fetch_moment_utc.isoformat()

    rows = fetch_rows(url, config.timeout_seconds)
    star_field_mismatches = ensure_consistent_star_fields(rows)
    canonical_payload = build_canonical_payload(rows, config.hostname, fetch_moment_utc)

    raw_snapshot = {
        "source": {
            "provider": "NASA Exoplanet Archive",
            "table": "pscomppars",
            "hostname": config.hostname,
            "tap_sync_url": NASA_TAP_SYNC_URL,
            "query": query,
            "request_url": url,
        },
        "fetched_at_utc": fetched_at_utc,
        "row_count": len(rows),
        "rows": rows,
    }

    output_dir = config.output_dir
    write_json(output_dir / "raw-pscomppars.json", raw_snapshot)
    write_json(output_dir / "canonical-system.json", canonical_payload)

    print(f"Wrote NASA snapshot and canonical system payload to {output_dir}")
    if star_field_mismatches:
        print(
            "Warning: repeated star/system fields differed across planet rows: "
            + ", ".join(star_field_mismatches),
            file=sys.stderr,
        )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())