"""SGP4 propagation backend.

Ported from the sibling ``yksa_tdas`` service's ``location/propagate.py`` (pure
``sgp4`` + stdlib ``math``, no numpy). Builds a ``Satrec`` straight from the
stored canonical OMM dict via :func:`sgp4.omm.initialize` -- no TLE-line
round-trip -- propagates to the requested instant, and fills the TEME state plus
Earth-fixed (ECEF) and WGS84 geodetic views on the :class:`StateVector`.

The TEME->ECEF rotation uses GMST only (no polar motion / nutation); a future
Orekit backend will provide rigorous frames. That approximation is fine for the
sub-km display accuracy this service targets.
"""

from __future__ import annotations

import math
from datetime import datetime, timedelta, timezone

from sgp4 import omm as sgp4_omm
from sgp4.api import SGP4_ERRORS, Satrec, jday
from sgp4.propagation import gstime

from odm import PropagationError, PropagatorBackend, StateVector, parse_omm_epoch

WGS84_A_KM = 6378.137
WGS84_F = 1.0 / 298.257223563
WGS84_E2 = WGS84_F * (2.0 - WGS84_F)
EARTH_ROT_RAD_S = 7.292115e-5


#: Frames the pure-Python backend can transform between (GMST rotation only).
_SGP4_FRAMES = ("TEME", "ITRF")


class Sgp4Backend(PropagatorBackend):
    name = "sgp4"
    # SGP4 natively yields TEME; the GMST rotation gives an Earth-fixed view.
    display_frames = ("TEME", "ITRF")

    def transform(self, states, frame_in, frame_out):
        frame_in = frame_in.upper()
        frame_out = frame_out.upper()
        if frame_in == frame_out:
            return list(states)
        if frame_in not in _SGP4_FRAMES or frame_out not in _SGP4_FRAMES:
            raise PropagationError(
                f"sgp4 backend only transforms between {_SGP4_FRAMES}; "
                f"got {frame_in!r}->{frame_out!r} (use the orekit backend for more)"
            )
        out = []
        for sv in states:
            if sv.epoch is None:
                raise PropagationError("state needs an epoch to rotate frames")
            jd_ut1 = _jd_of(sv.epoch)
            if frame_in == "TEME":  # TEME -> ITRF
                r, v = _teme_to_ecef(sv.r_km, sv.v_kms, jd_ut1)
            else:  # ITRF -> TEME
                r, v = _ecef_to_teme(sv.r_km, sv.v_kms, jd_ut1)
            out.append(_replace_frame(sv, frame_out, r, v))
        return out

    def state_at(self, omm: dict, at: datetime | None = None) -> StateVector:
        sat = _satrec_from_omm(omm)
        element_epoch = _epoch_of(sat)
        when = at if at is not None else element_epoch
        if when is None:
            raise PropagationError("OMM has no usable epoch and no target time given")
        return _state(sat, when, element_epoch)

    def ephemeris(
        self, omm: dict, start: datetime, stop: datetime, step_s: float,
    ) -> list[StateVector]:
        if step_s <= 0:
            raise PropagationError("step must be positive")
        if stop < start:
            raise PropagationError("stop must be on or after start")
        sat = _satrec_from_omm(omm)
        element_epoch = _epoch_of(sat)
        out: list[StateVector] = []
        t = start.astimezone(timezone.utc)
        stop = stop.astimezone(timezone.utc)
        step = timedelta(seconds=step_s)
        # Guard against runaway loops; callers cap this via OEM_MAX_POINTS.
        while t <= stop + timedelta(microseconds=1):
            out.append(_state(sat, t, element_epoch))
            t += step
        return out


def _satrec_from_omm(omm: dict) -> Satrec:
    sat = Satrec()
    sgp4_omm.initialize(sat, _normalise_for_sgp4(omm))
    return sat


def _normalise_for_sgp4(omm: dict) -> dict:
    """Return a copy whose EPOCH matches sgp4.omm's strict ``...%S.%f`` parser.

    Stored epochs occasionally lack fractional seconds (or carry a trailing
    ``Z``); sgp4.omm.initialize only accepts ``%Y-%m-%dT%H:%M:%S.%f``.
    """
    fields = dict(omm)
    epoch = fields.get("EPOCH")
    dt = parse_omm_epoch(epoch)
    if dt is not None:
        fields["EPOCH"] = dt.astimezone(timezone.utc).replace(tzinfo=None).strftime(
            "%Y-%m-%dT%H:%M:%S.%f"
        )
    return fields


def _epoch_of(sat: Satrec) -> datetime | None:
    try:
        jd = sat.jdsatepoch + sat.jdsatepochF
        unix = (jd - 2440587.5) * 86400.0
        return datetime.fromtimestamp(unix, tz=timezone.utc)
    except Exception:  # noqa: BLE001 -- defensive; a bad epoch just yields None
        return None


def _state(sat: Satrec, when: datetime, element_epoch: datetime | None) -> StateVector:
    when = when.astimezone(timezone.utc)
    jd, fr = jday(
        when.year, when.month, when.day,
        when.hour, when.minute, when.second + when.microsecond / 1e6,
    )
    err, r, v = sat.sgp4(jd, fr)
    if err != 0:
        raise PropagationError(SGP4_ERRORS.get(err, f"sgp4 error {err}"))

    r_ecef, v_ecef = _teme_to_ecef(r, v, jd + fr)
    geodetic = _ecef_to_geodetic(r_ecef)
    return StateVector(
        epoch=when,
        frame="TEME",
        r_km=(r[0], r[1], r[2]),
        v_kms=(v[0], v[1], v[2]),
        ecef_km=r_ecef,
        ecef_v_kms=v_ecef,
        geodetic=geodetic,
        element_epoch=element_epoch,
    )


def _jd_of(when: datetime) -> float:
    """Full Julian date (UT1≈UTC) for a UTC datetime."""
    when = when.astimezone(timezone.utc)
    jd, fr = jday(
        when.year, when.month, when.day,
        when.hour, when.minute, when.second + when.microsecond / 1e6,
    )
    return jd + fr


def _replace_frame(sv: StateVector, frame: str, r, v) -> StateVector:
    """Copy ``sv`` with a new frame label and position/velocity.

    When the target is the Earth-fixed frame, mirror the rotated state into the
    ECEF convenience fields so geodetic-dependent conversions keep working.
    """
    ecef_km = r if frame == "ITRF" else sv.ecef_km
    ecef_v_kms = v if frame == "ITRF" else sv.ecef_v_kms
    return StateVector(
        epoch=sv.epoch,
        frame=frame,
        r_km=tuple(r),
        v_kms=tuple(v),
        ecef_km=ecef_km,
        ecef_v_kms=ecef_v_kms,
        geodetic=sv.geodetic,
        element_epoch=sv.element_epoch,
        warnings=list(sv.warnings),
    )


def _teme_to_ecef(r, v, jd_ut1):
    """Rotate TEME position/velocity to Earth-fixed (ECEF) via GMST."""
    theta = gstime(jd_ut1)
    cos, sin = math.cos(theta), math.sin(theta)
    x = cos * r[0] + sin * r[1]
    y = -sin * r[0] + cos * r[1]
    z = r[2]
    vx = cos * v[0] + sin * v[1] + EARTH_ROT_RAD_S * y
    vy = -sin * v[0] + cos * v[1] - EARTH_ROT_RAD_S * x
    vz = v[2]
    return (x, y, z), (vx, vy, vz)


def _ecef_to_teme(r, v, jd_ut1):
    """Inverse of :func:`_teme_to_ecef`: Earth-fixed (ECEF) -> TEME via GMST."""
    theta = gstime(jd_ut1)
    cos, sin = math.cos(theta), math.sin(theta)
    x = cos * r[0] - sin * r[1]
    y = sin * r[0] + cos * r[1]
    z = r[2]
    # Undo the Earth-rotation term, then the rotation, to recover TEME velocity.
    a = v[0] - EARTH_ROT_RAD_S * r[1]
    b = v[1] + EARTH_ROT_RAD_S * r[0]
    vx = cos * a - sin * b
    vy = sin * a + cos * b
    vz = v[2]
    return (x, y, z), (vx, vy, vz)


def _ecef_to_geodetic(r) -> tuple[float, float, float]:
    """ECEF (km) -> WGS84 geodetic latitude/longitude (deg) and altitude (km)."""
    x, y, z = r
    lon = math.atan2(y, x)
    p = math.hypot(x, y)
    lat = math.atan2(z, p * (1.0 - WGS84_E2))
    alt = 0.0
    for _ in range(8):
        sin_lat = math.sin(lat)
        n = WGS84_A_KM / math.sqrt(1.0 - WGS84_E2 * sin_lat * sin_lat)
        alt = p / math.cos(lat) - n
        lat = math.atan2(z, p * (1.0 - WGS84_E2 * n / (n + alt)))
    return math.degrees(lat), math.degrees(lon), alt