use crate::ecs::components::{JointPalette, SkinnedMesh};
use crate::ecs::resources::SkeletonAssets;
use archetype_ecs::prelude::{Entity, EntityId, GlobalTransform, World};
use glam::Mat4;

/// System that computes the skinning palette (joint matrices)
///
/// For each bone i:
/// JointMatrix[i] = ModelInverse * WorldBoneMatrix[i] * InverseBindPose[i]
///
/// This system must run AFTER transform propagation (HierarchyUpdateSystem).
pub fn skinning_palette_system(world: &mut World) {
    let skeleton_assets = match world.resource::<SkeletonAssets>() {
        Some(assets) => assets.clone(),
        None => return,
    };

    // Step 1: Collect all skinned mesh entities to avoid borrow conflicts during iteration
    let mesh_entities: Vec<EntityId> = world
        .query::<(Entity, &SkinnedMesh)>()
        .iter()
        .map(|(e, _)| e)
        .collect();

    let mut updates = Vec::new();

    // Step 2: For each mesh, calculate the palette
    for entity in mesh_entities {
        let (skeleton_handle, bone_entities) = {
            if let Some(mesh) = world.get_component::<SkinnedMesh>(entity) {
                (mesh.skeleton, mesh.bone_entities.clone())
            } else {
                continue;
            }
        };

        let skeleton = match skeleton_assets.get(skeleton_handle) {
            Some(s) => s,
            None => continue,
        };

        // Get the model root's inverse global transform
        let model_inv = if let Some(global) = world.get_component::<GlobalTransform>(entity) {
            Mat4::from_scale_rotation_translation(global.scale, global.rotation, global.position)
                .inverse()
        } else {
            Mat4::IDENTITY
        };

        let mut matrices = Vec::with_capacity(skeleton.bones.len());

        for (i, bone) in skeleton.bones.iter().enumerate() {
            let joint_mat = if let Some(bone_entity) = bone_entities.get(i) {
                if let Some(global) = world.get_component::<GlobalTransform>(*bone_entity) {
                    let world_mat = Mat4::from_scale_rotation_translation(
                        global.scale,
                        global.rotation,
                        global.position,
                    );

                    model_inv * world_mat * bone.inverse_bind_pose
                } else {
                    Mat4::IDENTITY
                }
            } else {
                Mat4::IDENTITY
            };

            matrices.push(joint_mat);
        }

        updates.push((entity, matrices));
    }

    // Step 3: Apply updates
    for (entity, matrices) in updates {
        if let Some(palette) = world.get_component_mut::<JointPalette>(entity) {
            palette.matrices = matrices;
        } else {
            let _ = world.add_component(entity, JointPalette { matrices });
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::Skeleton;
    use crate::ecs::components::BoneJoint;
    use glam::{Quat, Vec3};

    #[test]
    fn test_skinning_palette_calculation_fixed() {
        let mut world = World::new();
        let mut skeleton_assets = SkeletonAssets::new();

        let mut skeleton = Skeleton::new();
        let ibp = Mat4::from_translation(Vec3::new(-1.0, 0.0, 0.0));
        skeleton.add_bone("Bone".to_string(), None, ibp);
        let skele_handle = skeleton_assets.add(skeleton);
        world.insert_resource(skeleton_assets);

        let bone_entity = world.spawn((
            BoneJoint {
                bone_index: 0,
                skeleton: skele_handle,
            },
            GlobalTransform {
                position: Vec3::new(2.0, 0.0, 0.0),
                rotation: Quat::IDENTITY,
                scale: Vec3::ONE,
            },
        ));

        let mesh_entity = world.spawn((
            SkinnedMesh {
                skeleton: skele_handle,
                mesh_handle: 0,
                bone_entities: vec![bone_entity],
                joint_offset: 0,
            },
            GlobalTransform::identity(),
        ));

        skinning_palette_system(&mut world);

        let palette = world.get_component::<JointPalette>(mesh_entity).unwrap();

        let expected_x = 1.0;
        let actual_x = palette.matrices[0].to_cols_array()[12];

        assert!((actual_x - expected_x).abs() < 0.001);
    }
}
