nightshade 0.8.2

use super::drag_math::{
    DragParams, apply_axis_drag, apply_free_drag, apply_plane_drag, compute_drag_sensitivity,
    determine_drag_mode,
};
use super::recreate_gizmo_for_mode;
use crate::editor::context::{self, EditorContext};
use crate::prelude::*;

const GIZMO_SCALE_FACTOR: f32 = 0.08;
const GIZMO_MIN_SCALE: f32 = 0.3;
const GIZMO_MAX_SCALE: f32 = 50.0;

fn set_gizmo_visible(world: &mut World, gizmo: &crate::ecs::gizmos::GizmoState, visible: bool) {
    if let Some(vis) = world.get_visibility_mut(gizmo.root) {
        vis.visible = visible;
    }

    let descendants = crate::ecs::transform::queries::query_descendants(world, gizmo.root);

    for descendant in descendants {
        if let Some(vis) = world.get_visibility_mut(descendant) {
            vis.visible = visible;
        }
    }
}

fn ensure_gizmo_exists(context: &mut EditorContext, world: &mut World) {
    use crate::ecs::gizmos;

    if context.gizmo_interaction.active.is_some() {
        return;
    }

    let dummy_entity = world.spawn_entities(
        LOCAL_TRANSFORM | LOCAL_TRANSFORM_DIRTY | GLOBAL_TRANSFORM,
        1,
    )[0];

    let gizmo = match context.gizmo_interaction.mode {
        crate::ecs::gizmos::GizmoMode::LocalTranslation
        | crate::ecs::gizmos::GizmoMode::GlobalTranslation => {
            gizmos::create_translation_gizmo(world, dummy_entity, context.gizmo_interaction.mode)
        }
        crate::ecs::gizmos::GizmoMode::Rotation => {
            gizmos::create_rotation_gizmo(world, dummy_entity)
        }
        crate::ecs::gizmos::GizmoMode::Scale => gizmos::create_scale_gizmo(world, dummy_entity),
    };

    world.add_components(gizmo.root, PARENT | IGNORE_PARENT_SCALE);
    world.update_parent(gizmo.root, None);
    despawn_recursive_immediate(world, dummy_entity);
    set_gizmo_visible(world, &gizmo, false);
    gizmos::set_gizmo_drag_visibility(world, &gizmo, crate::ecs::gizmos::GizmoDragMode::None);

    context.gizmo_interaction.active = Some(gizmo);
}

fn update_gizmo_parent_and_visibility(context: &mut EditorContext, world: &mut World) -> bool {
    use crate::ecs::gizmos;

    let Some(gizmo) = &context.gizmo_interaction.active else {
        return false;
    };

    let Some(selected_entity) = context.selection.primary() else {
        world.update_parent(gizmo.root, None);
        set_gizmo_visible(world, gizmo, false);
        context.gizmo_interaction.hover_axis = None;
        context.gizmo_interaction.drag_mode = crate::ecs::gizmos::GizmoDragMode::None;
        return false;
    };

    if world.resources.active_camera == Some(selected_entity) {
        world.update_parent(gizmo.root, None);
        set_gizmo_visible(world, gizmo, false);
        context.gizmo_interaction.hover_axis = None;
        context.gizmo_interaction.drag_mode = crate::ecs::gizmos::GizmoDragMode::None;
        return false;
    }

    let current_parent = world.get_parent(gizmo.root).and_then(|p| p.0);
    if current_parent != Some(selected_entity) {
        world.update_parent(gizmo.root, Some(Parent(Some(selected_entity))));
    }

    let modal_op = context.transform_edit.modal.operation;
    let in_modal_transform = modal_op != context::TransformOperation::None;

    if in_modal_transform {
        let desired_mode = match modal_op {
            context::TransformOperation::Grab => {
                Some(crate::ecs::gizmos::GizmoMode::LocalTranslation)
            }
            context::TransformOperation::Rotate => Some(crate::ecs::gizmos::GizmoMode::Rotation),
            context::TransformOperation::Scale => Some(crate::ecs::gizmos::GizmoMode::Scale),
            context::TransformOperation::None => None,
        };

        let needs_gizmo_switch = desired_mode.is_some_and(|desired| match desired {
            crate::ecs::gizmos::GizmoMode::LocalTranslation => !matches!(
                context.gizmo_interaction.mode,
                crate::ecs::gizmos::GizmoMode::LocalTranslation
                    | crate::ecs::gizmos::GizmoMode::GlobalTranslation
            ),
            _ => context.gizmo_interaction.mode != desired,
        });

        if needs_gizmo_switch {
            context.gizmo_interaction.mode = desired_mode.unwrap();
            recreate_gizmo_for_mode(context, world);
            return false;
        }

        set_gizmo_visible(world, gizmo, true);

        let drag_mode = match context.transform_edit.modal.constraint {
            context::TransformConstraint::None => crate::ecs::gizmos::GizmoDragMode::Free,
            context::TransformConstraint::X => crate::ecs::gizmos::GizmoDragMode::Axis(Vec3::x()),
            context::TransformConstraint::Y => crate::ecs::gizmos::GizmoDragMode::Axis(Vec3::y()),
            context::TransformConstraint::Z => crate::ecs::gizmos::GizmoDragMode::Axis(Vec3::z()),
        };
        gizmos::set_gizmo_drag_visibility(world, gizmo, drag_mode);

        let highlight_axis = match context.transform_edit.modal.constraint {
            context::TransformConstraint::None => None,
            context::TransformConstraint::X => Some(Vec3::x()),
            context::TransformConstraint::Y => Some(Vec3::y()),
            context::TransformConstraint::Z => Some(Vec3::z()),
        };
        gizmos::highlight_gizmo_axis(world, gizmo, highlight_axis);
    } else {
        set_gizmo_visible(world, gizmo, true);
        gizmos::set_gizmo_drag_visibility(world, gizmo, context.gizmo_interaction.drag_mode);
    }

    true
}

fn update_gizmo_scale(context: &mut EditorContext, world: &mut World) -> bool {
    let Some(gizmo) = &context.gizmo_interaction.active else {
        return false;
    };
    let Some(selected_entity) = context.selection.primary() else {
        return false;
    };

    if let Some(camera_entity) = world.resources.active_camera
        && let (Some(camera_transform), Some(target_transform)) = (
            world.get_global_transform(camera_entity),
            world.get_global_transform(selected_entity),
        )
    {
        let camera_pos = nalgebra_glm::vec3(
            camera_transform.0[(0, 3)],
            camera_transform.0[(1, 3)],
            camera_transform.0[(2, 3)],
        );
        let target_pos = nalgebra_glm::vec3(
            target_transform.0[(0, 3)],
            target_transform.0[(1, 3)],
            target_transform.0[(2, 3)],
        );
        let distance = nalgebra_glm::distance(&camera_pos, &target_pos);
        context.gizmo_interaction.camera_distance = distance;

        let scale = (distance * GIZMO_SCALE_FACTOR).clamp(GIZMO_MIN_SCALE, GIZMO_MAX_SCALE);

        if let Some(transform) = world.get_local_transform(gizmo.root) {
            let mut new_transform = *transform;
            new_transform.translation = Vec3::zeros();
            new_transform.scale = nalgebra_glm::vec3(scale, scale, scale);
            new_transform.rotation = Quat::identity();
            world.assign_local_transform(gizmo.root, new_transform);
        }
    }

    let in_modal = context.transform_edit.modal.operation != context::TransformOperation::None;
    !in_modal
}

fn end_drag(context: &mut EditorContext, world: &mut World) -> bool {
    use crate::ecs::gizmos;

    let initial_transforms = std::mem::take(&mut context.gizmo_interaction.drag_initial_transforms);
    let modified =
        context.commit_selection_transforms(world, initial_transforms, "Transform selection");
    context.gizmo_interaction.drag_mode = crate::ecs::gizmos::GizmoDragMode::None;
    if let Some(gizmo) = &context.gizmo_interaction.active {
        gizmos::set_gizmo_drag_visibility(world, gizmo, crate::ecs::gizmos::GizmoDragMode::None);
    }
    modified
}

fn update_hover_axis(context: &mut EditorContext, world: &mut World, mouse_pos: Vec2) {
    use crate::ecs::gizmos;
    use crate::ecs::picking::queries::{PickingOptions, pick_entities};

    let all_picked = pick_entities(world, mouse_pos, PickingOptions::default());
    let hover_axis = context.gizmo_interaction.active.as_ref().and_then(|gizmo| {
        all_picked
            .iter()
            .find_map(|result| gizmos::check_gizmo_axis_hit(world, gizmo, result.entity))
    });

    if hover_axis != context.gizmo_interaction.hover_axis {
        context.gizmo_interaction.hover_axis = hover_axis;
        if let Some(gizmo) = &context.gizmo_interaction.active {
            gizmos::highlight_gizmo_axis(world, gizmo, hover_axis);
        }
    }
}

fn begin_drag(context: &mut EditorContext, world: &mut World, axis: Vec3) {
    use crate::ecs::gizmos;

    let new_drag_mode = determine_drag_mode(axis);
    context.gizmo_interaction.drag_mode = new_drag_mode;
    if let Some(gizmo) = &context.gizmo_interaction.active {
        gizmos::set_gizmo_drag_visibility(world, gizmo, new_drag_mode);
    }

    context.gizmo_interaction.drag_initial_transforms = context.capture_selection_transforms(world);
    context.begin_selection_transform_tracking(world);
}

fn apply_drag(context: &EditorContext, world: &mut World) {
    let Some(camera_entity) = world.resources.active_camera else {
        return;
    };
    let mouse_delta = world.resources.input.mouse.raw_mouse_delta;
    let sensitivity = compute_drag_sensitivity(context.gizmo_interaction.camera_distance);
    let Some(camera_transform) = world.get_global_transform(camera_entity) else {
        return;
    };

    let (camera_right, camera_up, _) = super::drag_math::extract_normalized_axes(camera_transform);

    let drag_params = DragParams {
        mode: context.gizmo_interaction.mode,
        camera_right,
        camera_up,
        mouse_delta,
        sensitivity,
    };

    let drag_mode = context.gizmo_interaction.drag_mode;
    for &entity in context.selection.iter() {
        match &drag_mode {
            crate::ecs::gizmos::GizmoDragMode::Axis(axis) => {
                apply_axis_drag(world, entity, *axis, &drag_params);
            }
            crate::ecs::gizmos::GizmoDragMode::Plane(plane_normal) => {
                apply_plane_drag(world, entity, *plane_normal, &drag_params);
            }
            crate::ecs::gizmos::GizmoDragMode::Free => {
                apply_free_drag(world, entity, &drag_params);
            }
            crate::ecs::gizmos::GizmoDragMode::None => {}
        }
    }
}

fn handle_gizmo_interaction(context: &mut EditorContext, world: &mut World) -> bool {
    if context.gizmo_interaction.active.is_none() {
        return false;
    }
    if context.selection.primary().is_none() {
        return false;
    }

    if world.resources.user_interface.consumed_event {
        let mut modified = false;
        if context.gizmo_interaction.is_dragging() {
            modified = end_drag(context, world);
            context.gizmo_interaction.hover_axis = None;
        }
        return modified;
    }

    let mouse_pos = world.resources.input.mouse.position;
    let mouse_pressed = world
        .resources
        .input
        .mouse
        .state
        .contains(MouseState::LEFT_CLICKED);
    let mouse_just_pressed = world
        .resources
        .input
        .mouse
        .state
        .contains(MouseState::LEFT_JUST_PRESSED);

    let is_dragging = context.gizmo_interaction.is_dragging();
    let mouse_in_selected_viewport =
        crate::editor::camera_controls::is_mouse_in_selected_viewport(world);

    if !mouse_in_selected_viewport && !is_dragging {
        if context.gizmo_interaction.hover_axis.is_some() {
            context.gizmo_interaction.hover_axis = None;
            if let Some(gizmo) = &context.gizmo_interaction.active {
                crate::ecs::gizmos::highlight_gizmo_axis(world, gizmo, None);
            }
        }
        return false;
    }

    if !is_dragging {
        update_hover_axis(context, world, mouse_pos);
        if mouse_just_pressed && let Some(axis) = context.gizmo_interaction.hover_axis {
            begin_drag(context, world, axis);
        }
        false
    } else if mouse_pressed {
        apply_drag(context, world);
        false
    } else {
        let modified = end_drag(context, world);
        if mouse_in_selected_viewport {
            update_hover_axis(context, world, mouse_pos);
        } else {
            context.gizmo_interaction.hover_axis = None;
        }
        modified
    }
}

pub fn update_gizmo(context: &mut EditorContext, world: &mut World) -> bool {
    ensure_gizmo_exists(context, world);

    if !update_gizmo_parent_and_visibility(context, world) {
        return false;
    }

    if !update_gizmo_scale(context, world) {
        return false;
    }

    handle_gizmo_interaction(context, world)
}