use kiss3d::color::Color;
use obj::raw::object::Polygon;
use rapier_testbed3d::KeyCode;
use rapier_testbed3d::TestbedViewer;
use rapier3d::parry::bounding_volume;
use rapier3d::parry::transformation::voxelization::FillMode;
use rapier3d::prelude::*;
use std::fs::File;
use std::io::BufReader;

pub async fn run(viewer: &mut TestbedViewer) -> anyhow::Result<()> {
    /*
     * Voxel geometry type selection.
     */

    let settings = viewer.example_settings_mut();

    let falling_objects = settings.get_or_set_string(
        "Falling objects",
        5, // Defaults to Mixed.
        vec![
            "Ball".to_string(),
            "Cuboid".to_string(),
            "Cylinder".to_string(),
            "Cone".to_string(),
            "Capsule".to_string(),
            "Mixed".to_string(),
        ],
    );

    let voxel_size_y = settings.get_or_set_f32("Voxel size y", 1.0, 0.5..=2.0);
    let voxel_size = Vector::new(1.0, voxel_size_y, 1.0);
    let test_ccd = settings.get_or_set_bool("Test CCD", false);

    // TODO: give a better placement to the objs.
    // settings.get_or_set_bool("Load .obj", false);
    let load_obj = false;

    /*
     * World
     */
    let mut world = PhysicsWorld::new();

    /*
     * Create a bowl for the ground
     */
    /*
     * Create the convex decompositions.
     */
    if load_obj {
        let geoms = models();
        let ngeoms = geoms.len();
        let width = (ngeoms as f32).sqrt() as usize;
        let num_duplications = 1; // 4;
        let shift = 7.0f32;

        for (igeom, obj_path) in geoms.into_iter().enumerate() {
            let deltas = Pose::IDENTITY;

            let mut shapes = Vec::new();
            println!("Parsing and decomposing: {obj_path}");

            let input = BufReader::new(File::open(obj_path).unwrap());

            if let Ok(model) = obj::raw::parse_obj(input) {
                let mut vertices: Vec<Vector> = model
                    .positions
                    .iter()
                    .map(|v| Vector::new(v.0, v.1, v.2))
                    .collect();
                let indices: Vec<_> = model
                    .polygons
                    .into_iter()
                    .flat_map(|p| match p {
                        Polygon::P(idx) => idx.into_iter(),
                        Polygon::PT(idx) => {
                            Vec::from_iter(idx.into_iter().map(|i| i.0)).into_iter()
                        }
                        Polygon::PN(idx) => {
                            Vec::from_iter(idx.into_iter().map(|i| i.0)).into_iter()
                        }
                        Polygon::PTN(idx) => {
                            Vec::from_iter(idx.into_iter().map(|i| i.0)).into_iter()
                        }
                    })
                    .collect();

                // Compute the size of the model, to scale it and have similar size for everything.
                let aabb =
                    bounding_volume::details::point_cloud_aabb(&deltas, vertices.iter().copied());
                let center = aabb.center();
                let center_v = Vector::new(center.x, center.y, center.z);
                let diag = (aabb.maxs - aabb.mins).length();

                vertices
                    .iter_mut()
                    .for_each(|p| *p = (*p - center_v) * 6.0 / diag);

                let indices: Vec<_> = indices
                    .chunks(3)
                    .map(|idx| [idx[0] as u32, idx[1] as u32, idx[2] as u32])
                    .collect();

                let decomposed_shape =
                    SharedShape::voxelized_mesh(&vertices, &indices, 0.1, FillMode::default());

                shapes.push(decomposed_shape);

                for k in 1..num_duplications + 1 {
                    let x = (igeom % width) as f32 * shift - 3.0;
                    let y = (igeom / width) as f32 * shift + 4.0;
                    let z = k as f32 * shift - 3.0;

                    let body = RigidBodyBuilder::fixed().translation(Vector::new(x, y, z));
                    let mut shapes_iter = shapes.iter();
                    let first_shape = shapes_iter.next().unwrap();
                    let (handle, _) = world.insert(body, ColliderBuilder::new(first_shape.clone()));

                    for shape in shapes_iter {
                        let collider = ColliderBuilder::new(shape.clone());
                        world.insert_collider(collider, Some(handle));
                    }
                }
            }
        }
    }

    /*
     * Create a voxelized wavy floor.
     */
    let mut samples: Vec<Vector> = vec![];
    let n = 200;
    for i in 0..n {
        for j in 0..n {
            let y = (i as f32 / n as f32 * 10.0).sin().clamp(-0.8, 0.8)
                * (j as f32 / n as f32 * 10.0).cos().clamp(-0.8, 0.8)
                * 16.0;

            samples.push(Vector::new(i as f32, y * voxel_size_y, j as f32));

            if i == 0 || i == n - 1 || j == 0 || j == n - 1 {
                // Create walls so the object at the edge don't fall into the infinite void.
                for k in 0..4 {
                    samples.push(Vector::new(
                        i as f32,
                        (y + k as f32) * voxel_size_y,
                        j as f32,
                    ));
                }
            }
        }
    }
    let collider = ColliderBuilder::voxels_from_points(voxel_size, &samples).build();
    let floor_aabb = collider.compute_aabb();
    world.insert_collider(collider, None);

    /*
     * Some dynamic primitives.
     */
    let nik = 30;
    let extents = floor_aabb.extents() * 0.75;
    let margin = (floor_aabb.extents() - extents) / 2.0;
    let ball_radius = 0.5;
    for i in 0..nik {
        for j in 0..5 {
            for k in 0..nik {
                let mut rb = RigidBodyBuilder::dynamic().translation(Vector::new(
                    floor_aabb.mins.x + margin.x + i as f32 * extents.x / nik as f32,
                    floor_aabb.maxs.y + j as f32 * 2.0,
                    floor_aabb.mins.z + margin.z + k as f32 * extents.z / nik as f32,
                ));
                if test_ccd {
                    rb = rb.linvel(Vector::new(0.0, -1000.0, 0.0)).ccd_enabled(true);
                }
                let falling_objects = if falling_objects == 5 {
                    j % 5
                } else {
                    falling_objects
                };

                let co = match falling_objects {
                    0 => ColliderBuilder::ball(ball_radius),
                    1 => ColliderBuilder::cuboid(ball_radius, ball_radius, ball_radius),
                    2 => ColliderBuilder::cylinder(ball_radius, ball_radius),
                    3 => ColliderBuilder::cone(ball_radius, ball_radius),
                    4 => ColliderBuilder::capsule_y(ball_radius, ball_radius),
                    _ => unreachable!(),
                };
                world.insert(rb, co);
            }
        }
    }

    // Add callback for handling voxels edition, and highlighting the voxel
    // pointed at by the mouse. We spawn two fake colliders that don’t interact
    // with anything. They are used as gizmos to indicate where the ray hits on voxels
    // by highlighting the voxel and drawing a small ball at the intersection.
    let hit_indicator_handle = world.insert_collider(
        ColliderBuilder::ball(0.1).collision_groups(InteractionGroups::none()),
        None,
    );
    let hit_highlight_handle = world.insert_collider(
        ColliderBuilder::cuboid(0.51, 0.51, 0.51).collision_groups(InteractionGroups::none()),
        None,
    );
    viewer.set_initial_collider_color(hit_indicator_handle, Color::new(0.5, 0.5, 0.1, 1.0));
    viewer.set_initial_collider_color(hit_highlight_handle, Color::new(0.1, 0.5, 0.1, 1.0));

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

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

            let Some((mouse_orig, mouse_dir)) = viewer.mouse().ray else {
                continue;
            };

            let ray = Ray::new(
                Vector::new(mouse_orig.x, mouse_orig.y, mouse_orig.z),
                Vector::new(mouse_dir.x, mouse_dir.y, mouse_dir.z),
            );
            let filter = QueryFilter {
                predicate: Some(&|_, co: &Collider| co.shape().as_voxels().is_some()),
                ..Default::default()
            };
            let query_pipeline = world.broad_phase.as_query_pipeline(
                world.narrow_phase.query_dispatcher(),
                &world.bodies,
                &world.colliders,
                filter,
            );

            if let Some((handle, hit)) =
                query_pipeline.cast_ray_and_get_normal(&ray, Real::MAX, true)
            {
                // Highlight the voxel.
                let hit_collider = &world.colliders[handle];
                let hit_pos = hit_collider.position();
                let hit_local_normal = hit_pos.rotation.inverse() * hit.normal;
                let voxels = hit_collider.shape().as_voxels().unwrap();
                let FeatureId::Face(id) = hit.feature else {
                    unreachable!()
                };
                let voxel_key = voxels.voxel_at_flat_id(id).unwrap();
                let voxel_center_local = voxels.voxel_center(voxel_key);
                let voxel_center = hit_pos.rotation * voxel_center_local + hit_pos.translation;
                let voxel_size = voxels.voxel_size();
                let hit_highlight = world.colliders.get_mut(hit_highlight_handle).unwrap();
                hit_highlight.set_translation(voxel_center);
                hit_highlight
                    .shape_mut()
                    .as_cuboid_mut()
                    .unwrap()
                    .half_extents = voxel_size / 2.0 + Vector::splat(0.001);
                viewer.update_collider(hit_highlight_handle, &world);

                // Show the hit point.
                let hit_pt = ray.point_at(hit.time_of_impact);
                let hit_indicator = world.colliders.get_mut(hit_indicator_handle).unwrap();
                hit_indicator.set_translation(hit_pt);
                hit_indicator.shape_mut().as_ball_mut().unwrap().radius = voxel_size.length() / 3.5;
                viewer.update_collider(hit_indicator_handle, &world);

                // If a relevant key was pressed, edit the shape.
                if viewer.keys().pressed(KeyCode::Space) {
                    let removal_mode = viewer.keys().pressed(KeyCode::LShift);
                    let voxels = world
                        .colliders
                        .get_mut(handle)
                        .unwrap()
                        .shape_mut()
                        .as_voxels_mut()
                        .unwrap();
                    let mut affected_key = voxel_key;

                    if !removal_mode {
                        // Find index of max absolute value component
                        let abs_normal = hit_local_normal.abs();
                        let imax = if abs_normal.x >= abs_normal.y && abs_normal.x >= abs_normal.z {
                            0
                        } else if abs_normal.y >= abs_normal.z {
                            1
                        } else {
                            2
                        };
                        let normal_arr =
                            [hit_local_normal.x, hit_local_normal.y, hit_local_normal.z];
                        if normal_arr[imax] >= 0.0 {
                            affected_key[imax] += 1;
                        } else {
                            affected_key[imax] -= 1;
                        }
                    }

                    voxels.set_voxel(affected_key, !removal_mode);
                    viewer.update_collider(handle, &world);
                }
            } else {
                // When there is no hit, move the indicators behind the camera.
                let behind_camera = mouse_orig - mouse_dir * 1000.0;
                let behind_camera_vect =
                    Vector::new(behind_camera.x, behind_camera.y, behind_camera.z);
                let hit_indicator = world.colliders.get_mut(hit_indicator_handle).unwrap();
                hit_indicator.set_translation(behind_camera_vect);
                let hit_highlight = world.colliders.get_mut(hit_highlight_handle).unwrap();
                hit_highlight.set_translation(behind_camera_vect);
            }
        }
    }
    Ok(())
}

fn models() -> Vec<String> {
    vec![
        // "assets/3d/camel_decimated.obj".to_string(),
        "assets/3d/chair.obj".to_string(),
        "assets/3d/cup_decimated.obj".to_string(),
        "assets/3d/dilo_decimated.obj".to_string(),
        "assets/3d/feline_decimated.obj".to_string(),
        "assets/3d/genus3_decimated.obj".to_string(),
        "assets/3d/hand2_decimated.obj".to_string(),
        "assets/3d/hand_decimated.obj".to_string(),
        "assets/3d/hornbug.obj".to_string(),
        "assets/3d/octopus_decimated.obj".to_string(),
        "assets/3d/rabbit_decimated.obj".to_string(),
        // "assets/3d/rust_logo.obj".to_string(),
        "assets/3d/rust_logo_simplified.obj".to_string(),
        "assets/3d/screwdriver_decimated.obj".to_string(),
        "assets/3d/table.obj".to_string(),
        "assets/3d/tstTorusModel.obj".to_string(),
        // "assets/3d/tstTorusModel2.obj".to_string(),
        // "assets/3d/tstTorusModel3.obj".to_string(),
    ]
}
