shape_maker/src/undo.rs

use bevy::math::Vec3Swizzles;
use bevy::prelude::*;
use bevy_prototype_lyon::prelude::*;

use crate::{ShapeData, ImageData, CreateShape, PointSize};

#[derive(Debug, Clone)]
pub enum UndoItem {
    Base,
    CreateShape { entity: Entity },
    DeleteShape { transform: Transform, edges: Vec<Vec2>, name: String,  note: String, shape: CreateShape, dm: DrawMode, old_entity: Entity },
    MoveItem { shape: Entity, from: Vec2, to: Vec2 },
    MoveDot { shape_data: Entity, dot: usize, from: Vec2, to: Vec2 },
    Rotate { entity: Entity, from: f32, to: f32 },
    NameChange { entity: Entity, from: String },
    NoteChange { entity: Entity, from: String },
    ChangeZ { entity: Entity, from: f32, to: f32 },
    ReScale { entity: Entity, from: Vec2, to: Vec2 },
    CreateImage { entity: Entity },
    DeleteImage { id: ImageData, transform: Transform, handle: Handle<Image>, old_entity: Entity },
}
pub struct UndoStack {
    items: Vec<UndoItem>,
    current_action: usize,
    last_valid: usize,
}

impl Default for UndoStack {
    fn default() -> Self {
        Self { items: [UndoItem::Base].to_vec(), current_action: 0, last_valid: 0 }
    }
}
impl UndoStack {
    pub fn push(&mut self, item: UndoItem) {
        bevy::log::info!("PUSHING: {:?}", item);
        if self.current_action < self.items.len() - 1 {
            // We need to do a "semi push"
            self.current_action += 1;
            self.items[self.current_action] = item;
        }
        else {
            // We push a completly new item and might need more space!
            self.items.push(item);
            self.current_action = self.items.len() - 1;
        }
        self.last_valid = self.current_action;
    }
    /// Pop the last item in the stack
    pub fn pop(&mut self) -> Option<&mut UndoItem> {
        // If this happens we are clearly into invalid territory
        assert!(self.current_action < self.items.len());
        if self.current_action > 0 {
            let item = &mut self.items[self.current_action];
            self.current_action -= 1;
            Some(item)
        }
        else {
            None
        }
    }
    // If exists, "unpops" the last popped item
    pub fn unpop(&mut self) -> Option<&mut UndoItem> {
        if self.current_action < self.last_valid {
            self.current_action += 1;
            let item = &mut self.items[self.current_action];
            Some(item)
        }
        else {
            None
        }
    }
}

pub fn undo_redo_sys(
    mut coms: Commands,
    mut stack: ResMut<UndoStack>,
    keyboard: Res<Input<KeyCode>>,
    p_size: Res<PointSize>,
    mut shapes: Query<&mut ShapeData>,
    mut transforms: Query<&mut Transform>,
    global_transforms: Query<&GlobalTransform>,
    mut images: Query<(&mut ImageData, &Handle<Image>)>,
    mut paths: Query<&mut bevy_prototype_lyon::prelude::Path>,
    draw_modes: Query<&DrawMode>,
) {
    if keyboard.just_pressed(KeyCode::Z) && keyboard.pressed(KeyCode::LControl) {
        let item = if keyboard.pressed(KeyCode::LShift) { stack.unpop() } else { stack.pop() };
        if let Some(item) = item {
            match item {
                UndoItem::Base => bevy::log::error!("POP/UNPOP: Got UndoItem::Base as a result!"),
                UndoItem::CreateShape { entity } => {
                    let transform = transforms.get(*entity).map(|t| *t).unwrap_or_default();
                    let mut sd = shapes.get_mut(*entity).unwrap();
                    let mut name = String::new();
                    let mut note = String::new();
                    std::mem::swap(&mut name, &mut sd.name);
                    std::mem::swap(&mut note, &mut sd.note);
                    let edges = sd.edges.iter().map(|e| transforms.get(*e).unwrap().translation.xy()).collect::<Vec<Vec2>>();
                    let dm = *draw_modes.get(*entity).unwrap();
                    let shape = sd.shape;

                    coms.entity(*entity).despawn_recursive();
                    *item = UndoItem::DeleteShape { transform, edges, name, note, shape, dm, old_entity: *entity };
                
                },
                UndoItem::DeleteShape {
                    transform,
                    edges,
                    name,
                    note,
                    shape,
                    dm,
                    old_entity,
                } => {
                    let main_shape = spawn_main_shape(&mut coms, *dm, *transform, *shape, edges.clone());

                    let dots_dm = DrawMode::Fill(FillMode::color(Color::RED));
                    let dot_shape = shapes::Circle { radius: p_size.0, center: Vec2::ZERO };
                    let center = coms.spawn_bundle(GeometryBuilder::build_as(
                        &dot_shape,
                        dots_dm,
                        Transform::from_xyz(0.0, 0.0, 0.5)
                    )).id();
                    let edges = edges.iter().map(|v| {
                        coms.spawn_bundle(GeometryBuilder::build_as(
                            &dot_shape,
                            dots_dm,
                            Transform::from_translation(v.extend(0.5))
                        )).id()
                    }).collect::<Vec<Entity>>();

                    let mut nname = String::new();
                    let mut nnote = String::new();
                    std::mem::swap(&mut nname, name);
                    std::mem::swap(&mut nnote, note);

                    let mut main_shape = coms.entity(main_shape);

                    edges.iter().for_each(|e| { main_shape.add_child(*e); });
                    let sd = ShapeData { name: nname, shape: *shape, main_shape: main_shape.id(), edges, center: Some(center), note: nnote };
                    main_shape
                        .insert(sd)
                        .add_child(center);

                    let oe = *old_entity;
                    *item = UndoItem::CreateShape { entity: main_shape.id() };

                    let f = |ne: &mut Entity| { if *ne == oe { *ne = main_shape.id(); }};
                    for i in stack.items.iter_mut() {
                        match i {
                            UndoItem::CreateShape { entity } => f(entity),
                            UndoItem::MoveDot { shape_data, ..} => f(shape_data),
                            UndoItem::MoveItem { shape, .. } => f(shape),
                            UndoItem::Rotate { entity, .. } => f(entity),
                            UndoItem::NameChange { entity, .. } => f(entity),
                            UndoItem::NoteChange { entity, .. } => f(entity),
                            UndoItem::ChangeZ { entity, .. } => f(entity),
                            UndoItem::ReScale { entity, .. } => f(entity),
                            _ => {},
                        }
                    }
                },
                UndoItem::MoveItem { shape, from, to } => {
                    if let Ok(mut t) = transforms.get_mut(*shape) {
                        t.translation = from.extend(t.translation.z);
                    }
                    let t = *from;
                    *from = *to;
                    *to = t;
                },
                UndoItem::MoveDot { shape_data, dot, from, to } => {
                    let sd = shapes.get(*shape_data).unwrap();
                    let dot = sd.edges[*dot];
                    let gt = global_transforms.get(dot).unwrap();

                    let rot = gt.to_scale_rotation_translation().1;
                    let ang = rot.to_euler(EulerRot::XYZ).2;
                    let delta = Mat2::from_angle(-ang) * (*from - gt.translation().xy());
                    if let Ok(mut t) = transforms.get_mut(dot) {
                        t.translation += delta.extend(0.0);
                    }
                    // We need to recalculate the center, and update the points to be the new relevant point from the center
                    let center_offset = crate::modify::calc_shape_center_offset(&transforms, &*sd);
                    // Update each edge's offset, and then move the main shape's translation
                    for edge in sd.edges.iter() {
                        if let Ok(mut t) = transforms.get_mut(*edge) {
                            t.translation -= center_offset.extend(0.0);
                        }
                    }
                    if let Ok(mut t) = transforms.get_mut(sd.main_shape) {
                        t.translation += (Mat2::from_angle(ang) * center_offset).extend(0.0);
                    }

                    // Now we need to update the shape itself
                    crate::modify::update_main_shape(&mut paths, &transforms, &*sd);

                    let t = *from;
                    *from = *to;
                    *to = t;
                },
                UndoItem::Rotate { entity, from, to } => {
                    if let Ok(mut t) = transforms.get_mut(*entity) {
                        t.rotation = Quat::from_rotation_z(from.to_radians());
                    }
                    let t = *from;
                    *from = *to;
                    *to = t;
                },
                UndoItem::NameChange { entity, from } => {
                    if let Ok(mut sd) = shapes.get_mut(*entity) {
                        std::mem::swap(&mut sd.name, from);
                    }
                    else if let Ok((mut id, _)) = images.get_mut(*entity) {
                        std::mem::swap(&mut id.name, from);
                    }
                    // No need to update the item, so no need to replace
                },
                UndoItem::ChangeZ { entity, from, to } => {
                    if let Ok(mut t) = transforms.get_mut(*entity) {
                        t.translation.z = *from;
                    }
                    let t = *from;
                    *from = *to;
                    *to = t;
                },
                UndoItem::ReScale { entity, from, to } => {
                    if let Ok(mut t) = transforms.get_mut(*entity) {
                        t.scale = from.extend(1.0);
                    }
                    let t = *from;
                    *from = *to;
                    *to = t;
                },
                UndoItem::CreateImage { entity } => {
                    if let Ok((mut id, handle)) = images.get_mut(*entity) {
                        let transform = transforms.get(*entity).map(|t| *t).unwrap_or(Transform::default());
                        let mut nid = ImageData::default();
                        std::mem::swap(&mut *id, &mut nid);

                        coms.entity(*entity).despawn_recursive();
                        *item = UndoItem::DeleteImage { id: nid, transform, handle: handle.clone(), old_entity: *entity };
                    }
                },
                UndoItem::DeleteImage { transform, handle, id, old_entity } => {
                    let mut nid = ImageData::default();
                    std::mem::swap(id, &mut nid);

                    let e = coms.spawn_bundle(SpriteBundle {
                        texture: handle.clone(),
                        transform: *transform,
                        ..Default::default()
                    })
                    .insert(nid).id();

                    let oe = *old_entity;
                    *item = UndoItem::CreateImage { entity: e };
                    
                    // Update each item with the same entity id
                    let f = |ne: &mut Entity| { if *ne == oe { *ne = e; }};
                    for i in stack.items.iter_mut() {
                        match i {
                            UndoItem::CreateShape { entity } => f(entity),
                            UndoItem::MoveItem { shape, .. } => f(shape),
                            UndoItem::Rotate { entity, .. } => f(entity),
                            UndoItem::NameChange { entity, .. } => f(entity),
                            UndoItem::NoteChange { entity, .. } => f(entity),
                            UndoItem::ChangeZ { entity, .. } => f(entity),
                            UndoItem::ReScale { entity, .. } => f(entity),
                            _ => {},
                        }
                    }
                },
                UndoItem::NoteChange { entity, from } => {
                    if let Ok(mut sd) = shapes.get_mut(*entity) {
                        std::mem::swap(&mut sd.note, from);
                    }
                    else if let Ok((mut id, _)) = images.get_mut(*entity) {
                        std::mem::swap(&mut id.note, from);
                    }
                },
            }
        }
        else {
            bevy::log::info!("Nothing to pop/unpop!");
        }
    }
}

fn spawn_main_shape(coms: &mut Commands, dm: DrawMode, trans: Transform, shape: CreateShape, edges: Vec<Vec2>) -> Entity {
    match shape {
        CreateShape::Triangle => {
            coms.spawn_bundle(GeometryBuilder::build_as(
                &shapes::Polygon { points: edges, closed: true },
                dm,
                trans
            )).id()
        },
        CreateShape::Square => {
            coms.spawn_bundle(GeometryBuilder::build_as(
                &shapes::Rectangle { extents: (edges[0] - edges[1]).abs(), ..Default::default() },
                dm,
                trans
            )).id()
        },
        CreateShape::Circle => {
            coms.spawn_bundle(GeometryBuilder::build_as(
                &shapes::Circle { radius: edges[0].length(), center: Vec2::ZERO },
                dm,
                trans
            )).id()
        },
        CreateShape::Capsule => unimplemented!("I should prob get rid of this..."),
    }
}