use avian2d::{dynamics::{joints::EntityConstraint, solver::{joint_damping, schedule::SubstepSolverSystems, solver_body::{SolverBody, SolverBodyInertia}}}, math::RecipOrZero};

use crate::prelude::*;

pub fn damping_plugin(app: &mut App) {
    let substeps = app.get_schedule_mut(SubstepSchedule).expect("Add SubstepSchedule first");
    substeps.add_systems(
        module_joint_damping::<FixedJoint>
        .after(joint_damping::<DistanceJoint>) // DistanceJoint is the last damping in vanilla avian
        .in_set(SubstepSolverSystems::Damping)
    );
}

#[derive(Component, Debug, Clone, Copy)]
pub struct ModuleJointDamping {
    pub distance: f32,
    pub angular: f32,
}
fn module_joint_damping<T: Component + EntityConstraint<2>>(
    bodies: Query<(&mut SolverBody, &SolverBodyInertia, &Transform)>,
    joints: Query<(&T, &ModuleJointDamping)>,
    time: Res<Time>,
) {
    let delta_secs = time.delta_secs();

    let mut dummy_body1 = SolverBody::DUMMY;
    let mut dummy_body2 = SolverBody::DUMMY;

    for (joint, damping) in &joints {
        let entities = joint.entities();

        let (mut body1, mut inertia1, mut transform1) = (&mut dummy_body1, &SolverBodyInertia::DUMMY, Transform::default());
        let (mut body2, mut inertia2, mut transform2) = (&mut dummy_body2, &SolverBodyInertia::DUMMY, Transform::default());

        // Get the solver bodies for the two jointed entities.
        if let Ok((body, inertia, transform)) = unsafe { bodies.get_unchecked(entities[0]) } {
            body1 = body.into_inner();
            inertia1 = inertia;
            transform1 = *transform;
        }
        if let Ok((body, inertia, transform)) = unsafe { bodies.get_unchecked(entities[1]) } {
            body2 = body.into_inner();
            inertia2 = inertia;
            transform2 = *transform;
        }

        let delta_omega = (body2.angular_velocity - body1.angular_velocity)
            * (damping.angular * delta_secs).min(1.0);

        if !body1.flags.is_kinematic() {
            body1.angular_velocity += delta_omega;
        }
        if !body2.flags.is_kinematic() {
            body2.angular_velocity -= delta_omega;
        }

        let position1 = transform1.translation.truncate() + body1.delta_position;
        let position2 = transform2.translation.truncate() + body2.delta_position;

        let relative_position = position1 - position2;

        let parallel_velocity1 = body1.linear_velocity.project_onto(relative_position.normalize());
        let parallel_velocity2 = body2.linear_velocity.project_onto(relative_position.normalize());
        
        let delta_v = (parallel_velocity2 - parallel_velocity1)
            * (damping.distance * delta_secs).min(1.0);

        let w1 = inertia1.effective_inv_mass();
        let w2 = inertia2.effective_inv_mass();

        let p = delta_v * (w1 + w2).recip_or_zero();

        body1.linear_velocity += p * w1;
        body2.linear_velocity -= p * w2;
    }
}