sidereon

//! SGP4 orbit propagation via the in-house `sidereon_core::astro::sgp4` module
//! (pure-Rust port of Vallado SGP4, bit-exact to non-FMA Vallado at 0 ULP).

use crate::errors;
use rustler::Error;
use rustler::{Encoder, Env, NifResult, Term};
use sidereon_core::astro::forces::{
    DragParameters, SchwarzschildRelativity, SolarRadiationPressure, SpaceWeatherSource,
    SphericalHarmonicGravityConfig, ThirdBodyGravity, ZonalCoefficients, ZonalDegrees,
    ZonalGravity,
};
use sidereon_core::astro::propagator::{
    ForceModelComponents, ForceModelKind, IntegratorKind, IntegratorOptions, StatePropagator,
};
use sidereon_core::astro::sgp4::{ElementSet, JulianDate, OpsMode, Satellite};
use sidereon_core::astro::state::CartesianState;
use sidereon_core::astro::time::civil;
use sidereon_core::astro::tle::TleElements;

type DateTuple = (i32, i32, i32);
type TimeTuple = (i32, i32, i32, i32);
struct DatetimeComponents {
    year: i32,
    month: i32,
    day: i32,
    hour: i32,
    minute: i32,
    second: i32,
    microsecond: i32,
}

fn parse_datetime_tuple(term: Term) -> NifResult<DatetimeComponents> {
    let ((y, m, d), (h, min, s, us)): (DateTuple, TimeTuple) = term.decode()?;
    Ok(DatetimeComponents {
        year: y,
        month: m,
        day: d,
        hour: h,
        minute: min,
        second: s,
        microsecond: us,
    })
}

pub(crate) fn get_map_val<'a, T: rustler::Decoder<'a>>(
    env: Env<'a>,
    map: Term<'a>,
    key: &str,
) -> NifResult<T> {
    let atom = rustler::types::atom::Atom::from_str(env, key)?;
    let val = map.map_get(atom.to_term(env))?;
    val.decode()
}

/// Decode the SGP4 [`OpsMode`] selector the Elixir side threads through the
/// pass/visibility/look-angle NIFs as the atom `:afspc` or `:improved`.
///
/// The whole point is that a satellite's passes, visibility, and look angle are
/// built with the *same* opsmode, so the propagation underneath them is
/// consistent rather than silently forced to AFSPC by an element-set helper.
pub(crate) fn opsmode_from_term<'a>(env: Env<'a>, term: Term<'a>) -> NifResult<OpsMode> {
    let atom: rustler::Atom = term.decode()?;
    if atom == rustler::types::atom::Atom::from_str(env, "improved")? {
        Ok(OpsMode::Improved)
    } else if atom == rustler::types::atom::Atom::from_str(env, "afspc")? {
        Ok(OpsMode::Afspc)
    } else {
        Err(rustler::Error::BadArg)
    }
}

fn normalized_force_tokens(forces: &[String]) -> Vec<String> {
    forces
        .iter()
        .map(|force| force.trim().to_ascii_lowercase())
        .filter(|force| !force.is_empty())
        .collect()
}

fn has_force(tokens: &[String], aliases: &[&str]) -> bool {
    tokens
        .iter()
        .any(|force| aliases.iter().any(|alias| force == alias))
}

fn no_composite_perturbations(tokens: &[String]) -> bool {
    tokens.iter().all(|force| {
        matches!(
            force.as_str(),
            "twobody" | "two_body" | "two-body" | "j2" | "composite"
        )
    })
}

fn srp_from_tokens(tokens: &[String]) -> NifResult<Option<SolarRadiationPressure>> {
    for token in tokens {
        if let Some(rest) = token.strip_prefix("srp:") {
            let fields: Vec<&str> = rest.split(':').collect();
            if fields.len() != 2 {
                return Err(Error::Term(Box::new(
                    "srp token must be srp:cr:area_to_mass",
                )));
            }
            let cr = fields[0]
                .parse::<f64>()
                .map_err(|_| Error::Term(Box::new("invalid srp cr")))?;
            let area_to_mass = fields[1]
                .parse::<f64>()
                .map_err(|_| Error::Term(Box::new("invalid srp area_to_mass")))?;
            return SolarRadiationPressure::new(cr, area_to_mass)
                .map(Some)
                .map_err(crate::errors::invalid_input);
        }
    }
    Ok(None)
}

fn spherical_harmonic_from_tokens(
    tokens: &[String],
) -> NifResult<Option<SphericalHarmonicGravityConfig>> {
    for token in tokens {
        let (kind, rest) = if let Some(rest) = token.strip_prefix("geopotential:") {
            ("earth", rest)
        } else if let Some(rest) = token.strip_prefix("spherical_harmonic:") {
            ("earth", rest)
        } else if let Some(rest) = token.strip_prefix("spherical-harmonic:") {
            ("earth", rest)
        } else if let Some(rest) = token.strip_prefix("egm96:") {
            ("egm96", rest)
        } else if let Some(rest) = token.strip_prefix("earth_phase_b:") {
            ("earth", rest)
        } else if let Some(rest) = token.strip_prefix("phase_b:") {
            ("earth", rest)
        } else {
            continue;
        };

        let fields: Vec<&str> = rest.split(':').collect();
        if fields.len() != 2 {
            return Err(Error::Term(Box::new(
                "geopotential token must be geopotential:degree:order",
            )));
        }
        let degree = fields[0]
            .parse::<u16>()
            .map_err(|_| Error::Term(Box::new("invalid geopotential degree")))?;
        let order = fields[1]
            .parse::<u16>()
            .map_err(|_| Error::Term(Box::new("invalid geopotential order")))?;
        let config = if kind == "egm96" {
            SphericalHarmonicGravityConfig::egm96(degree, order)
        } else {
            SphericalHarmonicGravityConfig::earth(degree, order)
        }
        .map_err(crate::errors::invalid_input)?;
        return Ok(Some(config));
    }
    Ok(None)
}

fn j2_only_zonal() -> ZonalGravity {
    ZonalGravity::new(
        sidereon_core::astro::constants::MU_EARTH,
        sidereon_core::astro::constants::RE_EARTH,
        ZonalDegrees::J2_ONLY,
        ZonalCoefficients::default(),
    )
}

/// Decode the binding's force-token list into the 0.15 force model selector.
///
/// Existing callers can keep passing `["twobody"]` or `["twobody", "j2"]`.
/// New composable callers may pass aliases such as `:j2_j6`, `:third_body`,
/// `:sun`, `:moon`, `{:srp, cr, area_to_mass_m2_kg}`, `:relativity`,
/// `{:geopotential, degree, order}`, and `:earth_phase_a` through the Elixir
/// wrapper.
pub(crate) fn force_model_kind(forces: &[String]) -> NifResult<ForceModelKind> {
    let tokens = normalized_force_tokens(forces);
    let tokens = if tokens.is_empty() {
        vec!["twobody".to_string()]
    } else {
        tokens
    };
    let srp = srp_from_tokens(&tokens)?;
    let spherical_harmonic = spherical_harmonic_from_tokens(&tokens)?;

    if has_force(&tokens, &["earth_phase_a", "phase_a"]) {
        return Ok(ForceModelKind::earth_phase_a(srp));
    }
    if has_force(&tokens, &["earth_phase_b", "phase_b"]) && spherical_harmonic.is_none() {
        return ForceModelKind::earth_phase_b(
            sidereon_core::astro::forces::EGM96_EMBEDDED_MAX_DEGREE,
            sidereon_core::astro::forces::EGM96_EMBEDDED_MAX_ORDER,
            srp,
        )
        .map_err(crate::errors::invalid_input);
    }

    let wants_j2 = has_force(&tokens, &["j2"]);
    let wants_full_zonal = has_force(
        &tokens,
        &["j2_j6", "j2-j6", "zonal", "zonal_j2_j6", "zonals"],
    );
    let wants_sun = has_force(&tokens, &["sun", "third_body_sun"]);
    let wants_moon = has_force(&tokens, &["moon", "third_body_moon"]);
    let wants_third_body = has_force(&tokens, &["third_body", "third-body", "sun_moon"]);
    let wants_relativity = has_force(&tokens, &["relativity", "schwarzschild"]);
    let wants_composite = has_force(&tokens, &["composite"])
        || wants_full_zonal
        || wants_sun
        || wants_moon
        || wants_third_body
        || wants_relativity
        || spherical_harmonic.is_some()
        || srp.is_some();

    if wants_j2 && !wants_composite && no_composite_perturbations(&tokens) {
        return Ok(ForceModelKind::two_body_j2());
    }
    if !wants_composite && !wants_j2 {
        return Ok(ForceModelKind::two_body());
    }

    if wants_composite && wants_j2 && no_composite_perturbations(&tokens) {
        return Ok(ForceModelKind::two_body_j2());
    }

    let mut components = ForceModelComponents::earth_two_body();
    if wants_full_zonal {
        components = components.with_zonal(ZonalGravity::earth_j2_through_j6());
    } else if wants_j2 {
        components = components.with_zonal(j2_only_zonal());
    }
    if let Some(spherical_harmonic) = spherical_harmonic {
        components = components.with_spherical_harmonic(spherical_harmonic);
    }
    if wants_third_body || (wants_sun && wants_moon) {
        components = components.with_third_body(ThirdBodyGravity::default());
    } else if wants_sun {
        components = components.with_third_body(ThirdBodyGravity::sun());
    } else if wants_moon {
        components = components.with_third_body(ThirdBodyGravity::moon());
    }
    if let Some(srp) = srp {
        components = components.with_solar_radiation_pressure(srp);
    }
    if wants_relativity {
        components = components.with_relativity(SchwarzschildRelativity::default());
    }

    Ok(ForceModelKind::composite(components))
}

pub(crate) fn elements_from_map<'a>(env: Env<'a>, tle_map: Term<'a>) -> NifResult<ElementSet> {
    // The Elixir side carries the SGP4 mean elements plus the full four-digit
    // epoch year and fractional day (it derives the year from its own datetime,
    // so no two-digit-year pivot is reconstructed here). The core moved the
    // TLE-to-IR mapping (and the exact Vallado days2mdhms/jday epoch math) behind
    // `TleElements::to_element_set`, so build that public IR and convert through
    // the single canonical path.
    let elements = TleElements {
        // The map path carries no catalog identity; "0" decodes to catalog 0,
        // matching the pre-strict behavior for synthesized element sets.
        catalog_number: "0".to_string(),
        classification: String::new(),
        international_designator: String::new(),
        epoch_year: get_map_val(env, tle_map, "epoch_year")?,
        epoch_day_of_year: get_map_val(env, tle_map, "epochdays")?,
        mean_motion_dot: get_map_val(env, tle_map, "mean_motion_dot")?,
        mean_motion_double_dot: get_map_val(env, tle_map, "mean_motion_double_dot")?,
        bstar: get_map_val(env, tle_map, "bstar")?,
        ephemeris_type: 0,
        elset_number: 0,
        inclination_deg: get_map_val(env, tle_map, "inclination_deg")?,
        raan_deg: get_map_val(env, tle_map, "raan_deg")?,
        eccentricity: get_map_val(env, tle_map, "eccentricity")?,
        arg_perigee_deg: get_map_val(env, tle_map, "arg_perigee_deg")?,
        mean_anomaly_deg: get_map_val(env, tle_map, "mean_anomaly_deg")?,
        mean_motion: get_map_val(env, tle_map, "mean_motion")?,
        rev_number: 0,
    };

    elements.to_element_set().map_err(errors::invalid_input)
}

pub(crate) fn propagate_with_elements_impl<'a>(
    env: Env<'a>,
    tle_map: Term<'a>,
    datetime_tuple: Term<'a>,
) -> NifResult<Term<'a>> {
    let ok = rustler::types::atom::Atom::from_str(env, "ok")?;
    let error = rustler::types::atom::Atom::from_str(env, "error")?;

    let elements = elements_from_map(env, tle_map)?;

    let dt = parse_datetime_tuple(datetime_tuple)?;

    // AFSPC opsmode: matches the historical sidereon behavior (which used the
    // third-party sgp4 crate's `_afspc_compatibility_mode` functions). The
    // Skyfield reference values stored in oracle tests were calibrated to
    // AFSPC, so we preserve that mode for compatibility.
    let satellite = match Satellite::from_elements_with_opsmode(&elements, OpsMode::Afspc) {
        Ok(s) => s,
        Err(e) => return Ok((error, format!("SGP4 init: {e}")).encode(env)),
    };

    // Compute the target Julian Date from the supplied UTC components and let
    // Satellite::propagate_jd subtract the satrec's *cached* exact epoch
    // internally. This avoids any precision drift that would otherwise come
    // from computing the epoch JD a second time on this side.
    let target_jd = match utc_components_to_jd_split(&dt) {
        Some(jd) => jd,
        None => return Ok((error, "invalid datetime").encode(env)),
    };

    match satellite.propagate_jd(target_jd) {
        Ok(pred) => {
            let pos = (pred.position[0], pred.position[1], pred.position[2]);
            let vel = (pred.velocity[0], pred.velocity[1], pred.velocity[2]);
            Ok((ok, (pos, vel)).encode(env))
        }
        Err(e) => Ok((error, format!("SGP4 propagate: {e}")).encode(env)),
    }
}

/// Adaptive Dormand-Prince 5(4) numerical propagation of a raw ECI Cartesian
/// state over a relative span `dt_seconds`. Marshals the public options into
/// the core driver and returns `{:ok, {{rx, ry, rz}, {vx, vy, vz}}}` or
/// `{:error, atom}`.
pub(crate) fn propagate_dp54_impl(
    env: Env<'_>,
    position_km: (f64, f64, f64),
    velocity_km_s: (f64, f64, f64),
    dt_seconds: f64,
    forces: Vec<String>,
    abs_tol: f64,
    rel_tol: f64,
) -> NifResult<Term<'_>> {
    propagate_dp54_impl_with_drag(
        env,
        position_km,
        velocity_km_s,
        dt_seconds,
        forces,
        abs_tol,
        rel_tol,
        None,
    )
}

#[allow(clippy::too_many_arguments)]
pub(crate) fn propagate_dp54_impl_with_drag(
    env: Env<'_>,
    position_km: (f64, f64, f64),
    velocity_km_s: (f64, f64, f64),
    dt_seconds: f64,
    forces: Vec<String>,
    abs_tol: f64,
    rel_tol: f64,
    drag: Option<DragParameters>,
) -> NifResult<Term<'_>> {
    let ok = rustler::types::atom::Atom::from_str(env, "ok")?;
    let error = rustler::types::atom::Atom::from_str(env, "error")?;

    let options = IntegratorOptions {
        abs_tol,
        rel_tol,
        ..IntegratorOptions::default()
    };
    let propagator = StatePropagator {
        initial: CartesianState::new(
            0.0,
            [position_km.0, position_km.1, position_km.2],
            [velocity_km_s.0, velocity_km_s.1, velocity_km_s.2],
        ),
        force_model: force_model_kind(&forces)?,
        integrator: IntegratorKind::Dp54,
        options,
        drag,
        space_weather: None,
    };

    match propagator.ephemeris(&[dt_seconds]) {
        Ok(states) => {
            let state = states[0];
            let pos = state.position_array();
            let vel = state.velocity_array();
            let r = (pos[0], pos[1], pos[2]);
            let v = (vel[0], vel[1], vel[2]);
            Ok((ok, (r, v)).encode(env))
        }
        Err(_) => {
            let reason = rustler::types::atom::Atom::from_str(env, "propagation_failed")?;
            Ok((error, reason).encode(env))
        }
    }
}

#[allow(clippy::too_many_arguments)]
pub(crate) fn propagate_dp54_impl_with_drag_and_space_weather(
    env: Env<'_>,
    position_km: (f64, f64, f64),
    velocity_km_s: (f64, f64, f64),
    epoch_tdb_seconds: f64,
    dt_seconds: f64,
    forces: Vec<String>,
    abs_tol: f64,
    rel_tol: f64,
    drag: DragParameters,
    table: rustler::ResourceArc<crate::space_weather::SpaceWeatherTableResource>,
) -> NifResult<Term<'_>> {
    let ok = rustler::types::atom::Atom::from_str(env, "ok")?;
    let error = rustler::types::atom::Atom::from_str(env, "error")?;

    let options = IntegratorOptions {
        abs_tol,
        rel_tol,
        ..IntegratorOptions::default()
    };
    let propagator = StatePropagator {
        initial: CartesianState::new(
            epoch_tdb_seconds,
            [position_km.0, position_km.1, position_km.2],
            [velocity_km_s.0, velocity_km_s.1, velocity_km_s.2],
        ),
        force_model: force_model_kind(&forces)?,
        integrator: IntegratorKind::Dp54,
        options,
        drag: Some(drag),
        space_weather: Some(SpaceWeatherSource::Table(table.table.clone())),
    };

    match propagator.ephemeris(&[epoch_tdb_seconds + dt_seconds]) {
        Ok(states) => {
            let state = states[0];
            let pos = state.position_array();
            let vel = state.velocity_array();
            let r = (pos[0], pos[1], pos[2]);
            let v = (vel[0], vel[1], vel[2]);
            Ok((ok, (r, v)).encode(env))
        }
        Err(_) => {
            let reason = rustler::types::atom::Atom::from_str(env, "propagation_failed")?;
            Ok((error, reason).encode(env))
        }
    }
}

/// Build a split-form Julian date `(jd_whole, jd_fraction)` from a UTC calendar
/// tuple. Returns `None` for an out-of-range month/day.
///
/// The calendar arithmetic delegates to [`civil::split_julian_date`] (the single
/// core civil-time conversion); only the input range guard and the microsecond
/// fold into fractional seconds remain on this side.
fn utc_components_to_jd_split(dt: &DatetimeComponents) -> Option<JulianDate> {
    if dt.month < 1 || dt.month > 12 || dt.day < 1 || dt.day > 31 {
        return None;
    }
    let second = dt.second as f64 + dt.microsecond as f64 / 1_000_000.0;
    let (jd_whole, jd_fraction) =
        civil::split_julian_date(dt.year, dt.month, dt.day, dt.hour, dt.minute, second);
    Some(JulianDate(jd_whole, jd_fraction))
}