use avian3d::{math::Scalar, prelude::*};
use bevy::prelude::{
    App, Component, FixedUpdate, IntoScheduleConfigs, Plugin, ResMut, Resource, Startup, Update,
};

const DRAG_COEFF: Scalar = 0.00001;
const MAX_THRUST: Scalar = 900.0;

#[derive(Resource, Default, Debug)]
struct CyberLean {
    pub lean: Scalar,
}

#[derive(Debug, Resource)]
pub struct CatControllerSettings {
    pub kp: Scalar,
    pub kd: Scalar,
    pub ki: Scalar,
}

impl Default for CatControllerSettings {
    fn default() -> Self {
        Self {
            kp: 60.0,
            kd: 30.0,
            ki: 2.0,
        }
    }
}

#[derive(Component, Debug, Clone, Copy)]
pub struct CatControllerState {
    pub roll_integral: Scalar,
    pub roll_prev: Scalar,
    roll_limit: Scalar,
}

impl Default for CatControllerState {
    fn default() -> Self {
        Self {
            roll_integral: Default::default(),
            roll_prev: Default::default(),
            roll_limit: 1.5,
        }
    }
}

impl CatControllerState {
    pub fn update_roll(&mut self, error: Scalar, dt: Scalar) -> (Scalar, Scalar) {
        let lim = self.roll_limit;
        self.roll_integral = (self.roll_integral + (error * dt)).min(lim).max(-lim);
        let derivative = (error - self.roll_prev) / dt;
        self.roll_prev = error;
        (derivative, self.roll_integral)
    }
}

mod systems {
    use avian3d::{
        math::{Scalar, FRAC_PI_2, PI},
        prelude::{
            ComputedCenterOfMass, ExternalForce, ExternalTorque, Gravity, LinearVelocity,
            RayCaster, RayHits, RigidBodyQueryReadOnly,
        },
    };
    use bevy::{
        ecs::system::{Populated, Single},
        prelude::{
            ButtonInput, Color, Gizmos, GlobalTransform, KeyCode, Quat, Res, ResMut, Time,
            Transform, Vec, Vec3, With, Without,
        },
    };

    use super::{CatControllerSettings, CatControllerState, CyberLean, DRAG_COEFF, MAX_THRUST};
    use crate::{
        bike::{CyberBikeBody, CyberWheel, WheelConfig, WheelState, WHEEL_RADIUS},
        input::InputState,
    };

    const FRAC_PI_3: Scalar = PI / 3.0;
    const FRAC_PI_4: Scalar = FRAC_PI_2 / 2.0;

    fn rotate_point(pt: &Vec3, rot: &Quat) -> Vec3 {
        // thanks to https://danceswithcode.net/engineeringnotes/quaternions/quaternions.html
        let [x, y, z] = pt.normalize().to_array();
        let qpt = Quat::from_xyzw(x, y, z, 0.0);
        // p' = rot^-1 * qpt * rot
        let rot_qpt = rot.inverse() * qpt * *rot;

        // why does this need to be inverted???
        -Vec3::from_array([rot_qpt.x, rot_qpt.y, rot_qpt.z])
    }

    pub(super) fn calculate_lean(
        bike_state: Single<(&LinearVelocity, &Transform), With<CyberBikeBody>>,
        wheels: Populated<&GlobalTransform, With<CyberWheel>>,
        input: Res<InputState>,
        gravity: Res<Gravity>,
        mut lean: ResMut<CyberLean>,
    ) {
        let mut wheels = wheels.iter();
        let w1 = wheels.next().unwrap();
        let w2 = wheels.next().unwrap();
        let base = (w1.translation() - w2.translation()).length().abs();
        let (velocity, xform) = bike_state.into_inner();
        let vel = velocity.dot(*xform.forward());
        let v_squared = vel.powi(2);
        let steering_angle = yaw_to_angle(input.yaw);
        let radius = base / steering_angle.tan();
        let gravity = gravity.0.length();
        let v2_r = v_squared / radius;
        let tan_theta = (v2_r / gravity).clamp(-FRAC_PI_3, FRAC_PI_3);

        if tan_theta.is_normal() || tan_theta == 0.0 {
            lean.lean = tan_theta.atan().clamp(-FRAC_PI_3, FRAC_PI_3);
        } else if lean.lean.abs() > 0.01 {
            lean.lean *= 0.5;
        }
    }

    pub(super) fn apply_lean(
        bike_query: Single<(&Transform, &mut ExternalTorque, &mut CatControllerState)>,
        wheels: Populated<(&WheelState, &GlobalTransform, &CyberWheel)>,
        time: Res<Time>,
        settings: Res<CatControllerSettings>,
        lean: Res<CyberLean>,
        mut gizmos: Gizmos,
    ) {
        let (xform, mut torque, mut control_vars) = bike_query.into_inner();

        let mut factor = 1.0;
        let mut rxform = GlobalTransform::default();
        let mut fxform = GlobalTransform::default();
        for (wheel_state, xform, cwheel) in wheels.iter() {
            if wheel_state.contact_point.is_none() {
                factor -= 0.25;
            }
            match cwheel {
                CyberWheel::Front => {
                    fxform = *xform;
                }
                CyberWheel::Rear => {
                    rxform = *xform;
                }
            }
        }

        let tork_axis = (rxform.translation() - fxform.translation()).normalize();

        let world_up = Vec3::Y; //xform.translation.normalize();
        let rot = Quat::from_axis_angle(*xform.back(), lean.lean);
        let target_up = rotate_point(&world_up, &rot).normalize();

        // show which is up
        gizmos.arrow(
            xform.translation,
            xform.translation + *xform.up(),
            bevy::color::palettes::basic::YELLOW,
        );

        // show which is forward
        gizmos.arrow(
            *xform.forward() + xform.translation,
            2.5 * *xform.forward() + xform.translation,
            bevy::color::palettes::basic::PURPLE,
        );

        let bike_right = xform.right();

        let roll_error = bike_right.dot(target_up);
        let pitch_error = world_up.dot(*xform.back());

        // only try to correct roll if we're not totally vertical
        if pitch_error.abs() < 0.95 {
            let (derivative, integral) = control_vars.update_roll(roll_error, time.delta_secs());
            let mag =
                (settings.kp * roll_error) + (settings.ki * integral) + (settings.kd * derivative);
            if mag.is_finite() {
                //let lean_force = factor * mag * *xform.left();
                let tork = factor * mag * tork_axis;

                torque.apply_torque(tork);

                gizmos.arrow(
                    xform.translation + *xform.up(),
                    xform.translation + *xform.up() + mag * xform.left(),
                    Color::WHITE,
                );
            }
        }
    }

    fn yaw_to_angle(yaw: Scalar) -> Scalar {
        let v = yaw.powi(5) * FRAC_PI_4;
        if v.is_normal() {
            v
        } else {
            0.0
        }
    }

    pub(super) fn suspension(
        bike_body_query: Single<
            (&Transform, &ComputedCenterOfMass, &mut ExternalForce),
            With<CyberBikeBody>,
        >,
        mut wheel_mesh_query: Populated<
            (&mut Transform, &mut WheelState, &WheelConfig, &CyberWheel),
            Without<CyberBikeBody>,
        >,
        caster_query: Populated<(&RayCaster, &RayHits, &CyberWheel)>,
        time: Res<Time>,
        mut gizmos: Gizmos,
    ) {
        let dt = time.delta().as_secs_f64() as Scalar;

        let mut front_caster = &RayCaster::default();
        let mut rear_caster = &RayCaster::default();

        let mut front_hits = &RayHits::default();
        let mut rear_hits = &RayHits::default();

        for (caster, hits, wheel) in caster_query.iter() {
            match wheel {
                CyberWheel::Front => {
                    front_caster = caster;
                    front_hits = hits;
                }
                CyberWheel::Rear => {
                    rear_caster = caster;
                    rear_hits = hits;
                }
            }
        }

        let (bike_xform, com, mut bike_forces) = bike_body_query.into_inner();

        for (mut xform, mut state, config, wheel) in wheel_mesh_query.iter_mut() {
            let (caster, hits) = match wheel {
                CyberWheel::Front => (front_caster, front_hits),
                CyberWheel::Rear => (rear_caster, rear_hits),
            };

            if let Some(hit) = hits.iter().next() {
                let dist = hit.distance;
                let hit_point = caster.global_origin() + caster.global_direction() * dist;

                let displacement = config.rest_dist - dist;
                let damper_vel = (state.displacement - displacement) / dt;

                let mag = config.konstant * displacement - config.damping * damper_vel;
                let mag = mag.max(0.0);

                state.force_normal = mag;
                state.contact_normal = Some(hit.normal);
                state.contact_point = Some(hit_point);
                state.displacement = displacement;

                let cdir = caster.direction.as_vec3();
                xform.translation = config.attach + (cdir * (dist - WHEEL_RADIUS));

                let force = hit.normal * mag * dt * 100.0;
                gizmos.arrow(
                    hit_point,
                    hit_point + force,
                    Color::linear_rgb(1., 0.5, 0.2),
                );
                bike_forces.apply_force_at_point(force, hit_point, bike_xform.translation + com.0);
            } else {
                xform.translation = config.attach + (caster.direction.as_vec3() * config.rest_dist);
                state.reset();
            }
        }
    }

    pub(super) fn wheel_action(
        bike_query: Single<
            (
                &Transform,
                &LinearVelocity,
                &mut ExternalForce,
                RigidBodyQueryReadOnly,
            ),
            With<CyberBikeBody>,
        >,
        mut wheels: Populated<(&mut WheelState, &WheelConfig, &CyberWheel)>,
        time: Res<Time>,
        input: Res<InputState>,
        mut gizmos: Gizmos,
    ) {
        let (bike_xform, bike_vel, mut bike_force, bike_body) = bike_query.into_inner();
        let bike_vel = bike_vel.0;

        let dt = time.delta().as_secs_f64() as Scalar;
        let yaw_angle = -yaw_to_angle(input.yaw);

        for (mut state, config, wheel) in wheels.iter_mut() {
            if let Some(contact_point) = state.contact_point {
                // if contact_point is Some, we also have a normal.
                let normal = state.contact_normal.unwrap().normalize();
                let max_force_mag = state.force_normal * config.friction;

                let rot = Quat::from_axis_angle(normal, yaw_angle);
                let forward = normal.cross(*bike_xform.right());
                let thrust_mag = input.throttle * MAX_THRUST * dt;
                let (thrust_dir, thrust_force) = match wheel {
                    CyberWheel::Rear => (forward, thrust_mag * 0.5),
                    CyberWheel::Front => (rot * forward, thrust_mag * 0.5),
                };

                let thrust = thrust_force * thrust_dir;

                let friction_dir = match wheel {
                    CyberWheel::Front => normal.cross(thrust_dir),
                    CyberWheel::Rear => normal.cross(*bike_xform.forward()),
                };

                let vel = bike_body.velocity_at_point(contact_point - bike_xform.translation);
                let friction_factor = -0.025 * vel.dot(friction_dir);
                let friction = (friction_factor / dt) * friction_dir;

                let diff = bike_vel - vel;
                bevy::log::debug!("{wheel:?}: vel diff: {diff:?} ({})", diff.length(),);

                let mut force = thrust + friction;
                force *= dt * 50.0;
                let force_mag = force.length();
                if force_mag > max_force_mag {
                    state.sliding = true;
                    force = force.normalize_or_zero() * max_force_mag;
                } else {
                    state.sliding = false;
                }

                bike_force.apply_force_at_point(
                    force,
                    contact_point,
                    bike_xform.translation + bike_body.center_of_mass.0,
                );
                gizmos.arrow(
                    contact_point,
                    contact_point + friction,
                    Color::linear_rgb(1., 1., 0.2),
                );
                gizmos.arrow(
                    contact_point,
                    contact_point + thrust,
                    Color::linear_rgb(1., 0., 0.2),
                );
            }
        }
    }

    pub(super) fn drag(
        bike_query: Single<
            (
                &Transform,
                &LinearVelocity,
                &ComputedCenterOfMass,
                &mut ExternalForce,
            ),
            With<CyberBikeBody>,
        >,
        mut gizmos: Gizmos,
        time: Res<Time>,
    ) {
        let (xform, vel, com, mut force) = bike_query.into_inner();

        let dt = time.delta_secs();

        let speed = vel.length_squared();
        let dspeed = speed * 0.0000001;
        let dir = vel.normalize_or_zero();
        let drag = -dspeed * dt * DRAG_COEFF * dir;
        if speed > 0.1 {
            force.apply_force_at_point(
                drag,
                xform.translation - (0.2 * *xform.down()),
                xform.translation + com.0,
            );
        }

        bevy::log::debug!("speed: {}, drag force: {}", vel.length(), drag.length());

        gizmos.arrow(
            xform.translation,
            xform.translation + drag * 10.,
            Color::linear_rgb(1.0, 0.0, 0.0),
        );
    }

    pub(super) fn tweak(
        mut config: Populated<&mut WheelConfig>,
        mut keys: ResMut<ButtonInput<KeyCode>>,
    ) {
        let keyset: std::collections::HashSet<_> = keys.get_pressed().collect();
        let shifted = keyset.contains(&KeyCode::ShiftLeft) || keyset.contains(&KeyCode::ShiftRight);
        let config = config.iter_mut();
        for ref mut c in config {
            for key in &keyset {
                match key {
                    KeyCode::KeyS => {
                        if shifted {
                            c.konstant += 0.2;
                        } else {
                            c.konstant -= 0.2;
                        }
                        bevy::log::info!(c.konstant);
                    }
                    KeyCode::KeyD => {
                        if shifted {
                            c.damping += 0.1;
                        } else {
                            c.damping -= 0.1;
                        }
                        bevy::log::info!(c.damping);
                    }

                    _ => continue,
                }
            }
        }
        let released: Vec<_> = keys.get_just_released().copied().collect();
        for key in released {
            keys.clear_just_released(key);
        }
    }
}

use systems::{apply_lean, calculate_lean, drag, suspension, tweak, wheel_action};

pub struct CyberPhysicsPlugin;

impl Plugin for CyberPhysicsPlugin {
    fn build(&self, app: &mut App) {
        app.init_resource::<CatControllerSettings>()
            .init_resource::<CyberLean>()
            .add_plugins((PhysicsPlugins::default(), PhysicsDebugPlugin::default()))
            .insert_resource(SubstepCount(12))
            .add_systems(Startup, |mut gravity: ResMut<Gravity>| {
                gravity.0 *= 1.0;
            })
            .add_systems(
                FixedUpdate,
                (
                    calculate_lean,
                    apply_lean,
                    suspension,
                    wheel_action, /* drag */
                )
                    .chain(),
            )
            .add_systems(Update, tweak);
    }
}