use std::f64::consts::PI;
use bevy::math::DVec2;
use bevy_replicon::client::confirm_history::ConfirmHistory;
use bevy_replicon::prelude::Replicated;
use bevy_replicon::shared::replicon_tick::RepliconTick;
use crate::prelude::*;
use crate::shared::config::planet::PlanetSpring;

const CORRECTION_FACTOR: f64 = 0.1;
const SNAP_THRESHOLD: f64 = 500.0;

#[derive(Component, Default)]
struct LocalState {
    position: DVec2,
    linvel: DVec2,
    angvel: f64,
    rotation_cos: f64,
    rotation_sin: f64,
}

#[derive(Component, Default)]
struct CorrectionTarget {
    position: DVec2,
    linvel: DVec2,
    angvel: f64,
    rotation_cos: f64,
    rotation_sin: f64,
    last_tick: Option<RepliconTick>,
}

pub fn prediction_plugin(app: &mut App) {
    app
        .add_observer(init_prediction)
        .add_systems(PreUpdate, (
            save_local_state.before(bevy_replicon::client::ClientSystems::Receive),
            record_server_correction.after(bevy_replicon::client::ClientSystems::Receive),
        ))
        .add_systems(FixedUpdate, apply_correction
            .after(PhysicsSystems::Prepare)
            .before(PhysicsSystems::StepSimulation));
}

fn init_prediction(
    trigger: On<Add, Replicated>,
    query: Query<(&Position, &Rotation, &LinearVelocity, Option<&AngularVelocity>), Without<PlanetSpring>>,
    mut commands: Commands,
) {
    let entity = trigger.event_target();
    let Ok((pos, rot, linvel, angvel)) = query.get(entity) else { return };
    commands.entity(entity).insert((
        LocalState {
            position: pos.0,
            linvel: linvel.0,
            angvel: angvel.map(|a| a.0).unwrap_or(0.0),
            rotation_cos: rot.cos,
            rotation_sin: rot.sin,
        },
        CorrectionTarget::default(),
    ));
}

fn save_local_state(
    mut query: Query<(&Position, &Rotation, &LinearVelocity, &AngularVelocity, &mut LocalState)>,
) {
    for (pos, rot, linvel, angvel, mut local) in &mut query {
        if pos.0.is_nan() { continue; }
        local.position = pos.0;
        local.linvel = linvel.0;
        local.angvel = angvel.0;
        local.rotation_cos = rot.cos;
        local.rotation_sin = rot.sin;
    }
}

fn record_server_correction(
    mut query: Query<(
        &mut Position, &mut Rotation, &mut LinearVelocity, &mut AngularVelocity,
        &ConfirmHistory,
        &mut CorrectionTarget, &LocalState,
    )>,
) {
    for (mut pos, mut rot, mut linvel, mut angvel, history, mut target, local) in &mut query {
        let tick = history.last_tick();
        if target.last_tick != Some(tick) {
            target.last_tick = Some(tick);
            target.position = pos.0;
            target.linvel = linvel.0;
            target.angvel = angvel.0;
            target.rotation_cos = rot.cos;
            target.rotation_sin = rot.sin;
        }
        pos.0 = local.position;
        rot.cos = local.rotation_cos;
        rot.sin = local.rotation_sin;
        linvel.0 = local.linvel;
        angvel.0 = local.angvel;
    }
}

fn apply_correction(
    mut query: Query<(&mut Position, &mut Rotation, &mut LinearVelocity, &mut AngularVelocity, &CorrectionTarget)>,
) {
    for (mut pos, mut rot, mut linvel, mut angvel, target) in &mut query {
        let Some(_) = target.last_tick else { continue };

        let pos_err = target.position - pos.0;
        if pos_err.length() > SNAP_THRESHOLD {
            pos.0 = target.position;
            linvel.0 = target.linvel;
            angvel.0 = target.angvel;
            rot.cos = target.rotation_cos;
            rot.sin = target.rotation_sin;
        } else {
            pos.0 += pos_err * CORRECTION_FACTOR;
            let cur_linvel = linvel.0;
            linvel.0 += (target.linvel - cur_linvel) * CORRECTION_FACTOR;
            let cur_angvel = angvel.0;
            angvel.0 += (target.angvel - cur_angvel) * CORRECTION_FACTOR;

            let angle_local = rot.sin.atan2(rot.cos);
            let angle_target = target.rotation_sin.atan2(target.rotation_cos);
            let diff = ((angle_target - angle_local + PI).rem_euclid(2.0 * PI)) - PI;
            let new_angle = angle_local + diff * CORRECTION_FACTOR;
            rot.cos = new_angle.cos();
            rot.sin = new_angle.sin();
        }
    }
}