use rapier_testbed3d::{KeyCode, TestbedViewer};
use rapier3d::prelude::*;

pub async fn run(viewer: &mut TestbedViewer) -> anyhow::Result<()> {
    /*
     * World
     */
    let mut world = PhysicsWorld::new();

    /*
     * Ground.
     */
    let ground_size = Vector::new(60.0, 0.4, 60.0);
    let nsubdivs = 100;

    let heights = Array2::from_fn(nsubdivs + 1, nsubdivs + 1, |i, j| {
        -(i as f32 * ground_size.x / (nsubdivs as f32) / 2.0).cos()
            - (j as f32 * ground_size.z / (nsubdivs as f32) / 2.0).cos()
    });

    let collider = ColliderBuilder::heightfield(heights, ground_size)
        .translation(Vector::new(-7.0, 0.0, 0.0))
        .friction(1.0);
    world.insert_collider(collider, None);

    /*
     * Vehicle we will control manually, simulated using joints.
     * Strongly inspired from https://github.com/h3r2tic/cornell-mcray/blob/main/src/car.rs (MIT/Apache license).
     */
    const CAR_GROUP: Group = Group::GROUP_1;

    let wheel_params = [
        Vector::new(0.6874, 0.2783, -0.7802),
        Vector::new(-0.6874, 0.2783, -0.7802),
        Vector::new(0.64, 0.2783, 1.0254),
        Vector::new(-0.64, 0.2783, 1.0254),
    ];
    // TODO: lower center of mass?
    // let mut center_of_mass = wheel_params.iter().sum().unwrap() / 4.0;
    // center_of_mass.y = 0.0;

    let suspension_height = 0.12;
    let max_steering_angle = 35.0f32.to_radians();
    let drive_strength = 1.0;
    let wheel_radius = 0.28;
    let car_position = Vector::new(0.0, wheel_radius + suspension_height, 0.0);
    let body_position_in_car_space = Vector::new(0.0, 0.4739, 0.0);

    let body_position = car_position + body_position_in_car_space;

    let body_co = ColliderBuilder::cuboid(0.65, 0.3, 0.9)
        .density(100.0)
        .collision_groups(InteractionGroups::new(
            CAR_GROUP,
            !CAR_GROUP,
            InteractionTestMode::And,
        ));
    let body_rb = RigidBodyBuilder::dynamic()
        .translation(body_position)
        .build();
    let (body_handle, _) = world.insert(body_rb, body_co);

    let mut steering_joints = vec![];
    let mut motor_joints = vec![];

    for (wheel_id, wheel_pos_in_car_space) in wheel_params.into_iter().enumerate() {
        let is_front = wheel_id >= 2;
        let wheel_center: Vector = car_position + wheel_pos_in_car_space;

        let axle_mass_props = MassProperties::from_ball(100.0, wheel_radius);
        let axle_rb = RigidBodyBuilder::dynamic()
            .translation(wheel_center)
            .additional_mass_properties(axle_mass_props);
        let axle_handle = world.insert_body(axle_rb);

        // This is a fake cylinder collider that we add only because our testbed can
        // only render colliders. Setting it as sensor makes it show up as wireframe.
        let wheel_fake_co = ColliderBuilder::cylinder(wheel_radius / 2.0, wheel_radius)
            .rotation(Vector::Z * std::f32::consts::FRAC_PI_2)
            .sensor(true)
            .density(0.0)
            .collision_groups(InteractionGroups::none());

        // The actual wheel collider. Simulating the wheel as a ball is interesting as it
        // is mathematically simpler than a cylinder and cheaper to compute for collision-detection.
        let wheel_co = ColliderBuilder::ball(wheel_radius)
            .density(100.0)
            .collision_groups(InteractionGroups::new(
                CAR_GROUP,
                !CAR_GROUP,
                InteractionTestMode::And,
            ))
            .friction(1.0);
        let wheel_rb = RigidBodyBuilder::dynamic().translation(wheel_center);
        let (wheel_handle, _) = world.insert(wheel_rb, wheel_co);
        world.insert_collider(wheel_fake_co, Some(wheel_handle));

        let suspension_attachment_in_body_space =
            wheel_pos_in_car_space - body_position_in_car_space;

        // Suspension between the body and the axle.
        let mut locked_axes = JointAxesMask::LIN_X
            | JointAxesMask::LIN_Z
            | JointAxesMask::ANG_X
            | JointAxesMask::ANG_Z;
        if !is_front {
            locked_axes |= JointAxesMask::ANG_Y;
        }

        let mut suspension_joint = GenericJointBuilder::new(locked_axes)
            .limits(JointAxis::LinY, [0.0, suspension_height])
            .motor_position(JointAxis::LinY, 0.0, 1.0e4, 1.0e3)
            .local_anchor1(suspension_attachment_in_body_space);

        if is_front {
            suspension_joint =
                suspension_joint.limits(JointAxis::AngY, [-max_steering_angle, max_steering_angle]);
        }

        let body_axle_joint_handle =
            world.insert_impulse_joint(body_handle, axle_handle, suspension_joint);

        // Joint between the axle and the wheel.
        let wheel_joint = RevoluteJointBuilder::new(Vector::X);
        let wheel_joint_handle = world.insert_impulse_joint(axle_handle, wheel_handle, wheel_joint);

        if is_front {
            steering_joints.push(body_axle_joint_handle);
            motor_joints.push(wheel_joint_handle);
        }
    }

    /*
     * Create some cubes on the ground.
     */
    // let num = 8;
    // let rad = 0.1;
    //
    // let shift = rad * 2.0;
    // let centerx = shift * (num / 2) as f32;
    // let centery = rad;
    //
    // for j in 0usize..1 {
    //     for k in 0usize..4 {
    //         for i in 0..num {
    //             let x = i as f32 * shift - centerx;
    //             let y = j as f32 * shift + centery;
    //             let z = k as f32 * shift + centerx;
    //
    //             let rigid_body = RigidBodyBuilder::dynamic().translation(Vector::new(x, y, z));
    //             let handle = bodies.insert(rigid_body);
    //             let collider = ColliderBuilder::cuboid(rad, rad, rad);
    //             colliders.insert_with_parent(collider, handle, &mut bodies);
    //         }
    //     }
    // }

    /*
     * Set up the testbed.
     */
    viewer.set_world(&mut world);
    viewer.look_at(Vec3::new(10.0, 10.0, 10.0), Vec3::ZERO);

    while viewer.render_frame(&mut world).await {
        if viewer.simulating() {
            world.step();

            /*
             * Control the wheels motors.
             */
            let mut thrust = 0.0;
            let mut steering = 0.0;
            let mut boost = 1.0;

            for key in viewer.keys().get_pressed() {
                match *key {
                    KeyCode::Right => {
                        steering = -1.0;
                    }
                    KeyCode::Left => {
                        steering = 1.0;
                    }
                    KeyCode::Up => {
                        thrust = -drive_strength;
                    }
                    KeyCode::Down => {
                        thrust = drive_strength;
                    }
                    KeyCode::RShift => {
                        boost = 1.5;
                    }
                    _ => {}
                }
            }
            let mut should_wake_up = false;
            if thrust != 0.0 || steering != 0.0 {
                should_wake_up = true;
            }

            // Apply steering to the axles.
            for steering_handle in &steering_joints {
                let steering_joint = world
                    .impulse_joints
                    .get_mut(*steering_handle, should_wake_up)
                    .unwrap();
                steering_joint.data.set_motor_position(
                    JointAxis::AngY,
                    max_steering_angle * steering,
                    1.0e4,
                    1.0e3,
                );
            }

            // Apply thrust.
            // Pseudo-differential adjusting speed of engines depending on steering arc
            // Higher values result in more drifty behavior.
            let differential_strength = 0.5;
            let sideways_shift = (max_steering_angle * steering).sin() * differential_strength;
            let speed_diff = if sideways_shift > 0.0 {
                f32::hypot(1.0, sideways_shift)
            } else {
                1.0 / f32::hypot(1.0, sideways_shift)
            };

            let ms = [1.0 / speed_diff, speed_diff];
            for (motor_handle, &ms) in motor_joints.iter().copied().zip(ms.iter()) {
                let motor_joint = world
                    .impulse_joints
                    .get_mut(motor_handle, should_wake_up)
                    .unwrap();
                motor_joint.data.set_motor_velocity(
                    JointAxis::AngX,
                    -30.0 * thrust * ms * boost,
                    1.0e2,
                );
            }
        }
    }
    Ok(())
}
