use avian3d::prelude::*;
use bevy::prelude::*;

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

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

impl Default for CatControllerSettings {
    fn default() -> Self {
        Self {
            kp: 15.0,
            kd: 1.5,
            ki: 0.5,
        }
    }
}

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

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

impl CatControllerState {
    pub fn update_roll(&mut self, error: f32, dt: f32) -> (f32, f32) {
        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 std::f32::consts::{FRAC_PI_3, FRAC_PI_4};

    use avian3d::prelude::*;
    use bevy::prelude::*;

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

    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: Query<(&LinearVelocity, &Transform), With<CyberBikeBody>>,
        wheels: Query<&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.single();
        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() {
            lean.lean = tan_theta.atan().clamp(-FRAC_PI_3, FRAC_PI_3);
        } else {
            //lean.lean = 0.0;
        }
    }

    pub(super) fn apply_lean(
        mut bike_query: Query<(&Transform, &mut ExternalForce, &mut CatControllerState)>,
        wheels: Query<&WheelState>,
        time: Res<Time>,
        settings: Res<CatControllerSettings>,
        lean: Res<CyberLean>,
        mut gizmos: Gizmos,
    ) {
        let (xform, mut force, mut control_vars) = bike_query.single_mut();

        let mut factor = 1.0;
        for wheel in wheels.iter() {
            if wheel.contact_point.is_none() {
                factor -= 0.25;
            }
        }

        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();
                force.apply_force_at_point(
                    lean_force,
                    xform.translation + *xform.up(),
                    xform.translation,
                );
                gizmos.arrow(
                    xform.translation + *xform.up(),
                    xform.translation + *xform.up() + lean_force,
                    Color::WHITE,
                );
            }
        }
    }

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

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

        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, mut bike_forces) = bike_body_query.single_mut();

        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,
                    caster.global_origin(),
                    bike_xform.translation,
                );
            } else {
                xform.translation = config.attach + (caster.direction.as_vec3() * config.rest_dist);
                state.reset();
            }
        }
    }

    pub(super) fn steering(
        mut bike_query: Query<
            (
                &Transform,
                &LinearVelocity,
                &mut ExternalForce,
                RigidBodyQueryReadOnly,
            ),
            With<CyberBikeBody>,
        >,
        mut wheels: Query<(&mut WheelState, &WheelConfig, &CyberWheel)>,
        time: Res<Time>,
        input: Res<InputState>,
        mut gizmos: Gizmos,
    ) {
        let Ok((bike_xform, bike_vel, mut bike_force, bike_body)) = bike_query.get_single_mut()
        else {
            bevy::log::warn!("No entity for bike_query.");
            return;
        };
        let bike_vel = bike_vel.0;

        let dt = time.delta().as_secs_f32();
        let max_thrust = 1000.0;
        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_dir, thrust_force) = match wheel {
                    CyberWheel::Rear => (forward, input.throttle * max_thrust * dt),
                    CyberWheel::Front => (rot * forward, 0.0),
                };

                let thrust = thrust_force * thrust_dir;

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

                // bevy::log::debug!(
                //     "{wheel:?}, thrust_dir: {thrust_dir:?}, friction_dir: {friction_dir:?}, dot: {}",
                //     thrust_dir.dot(friction_dir)
                // );

                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;

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

                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;
                    //bevy::log::info!("clamping {wheel:?} to {max_force_mag}: {friction_factor}");
                } else {
                    state.sliding = false;
                }

                bike_force.apply_force_at_point(force, contact_point, bike_xform.translation);
                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 tweak(
        mut config: Query<&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, steering, suspension, tweak};

pub struct CyberPhysicsPlugin;

impl Plugin for CyberPhysicsPlugin {
    fn build(&self, app: &mut App) {
        app.init_resource::<CatControllerSettings>()
            .register_type::<CatControllerSettings>()
            .register_type::<CatControllerState>()
            .init_resource::<CyberLean>()
            .register_type::<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, steering).chain(),
            )
            .add_systems(Update, tweak);
    }
}