common/

undo_redo.rs

use crate::event::{Event, EventHub, Origin, UndoRedoEvent};
use crate::types::EntityId;
use anyhow::{Result, anyhow};
use std::any::Any;
use std::collections::HashMap;
use std::sync::Arc;

/// Trait for commands that can be undone and redone.
///
/// Implementors can optionally support command merging by overriding the
/// `can_merge` and `merge` methods. This allows the UndoRedoManager to combine
/// multiple commands of the same type into a single command, which is useful for
/// operations like continuous typing or dragging.
pub trait UndoRedoCommand: Send {
    /// Undoes the command, reverting its effects
    fn undo(&mut self) -> Result<()>;

    /// Redoes the command, reapplying its effects
    fn redo(&mut self) -> Result<()>;

    /// Returns true if this command can be merged with the other command.
    ///
    /// By default, commands cannot be merged. Override this method to enable
    /// merging for specific command types.
    ///
    /// # Example
    /// ```test
    /// fn can_merge(&self, other: &dyn UndoRedoCommand) -> bool {
    ///     // Check if the other command is of the same type
    ///     if let Some(_) = other.as_any().downcast_ref::<Self>() {
    ///         return true;
    ///     }
    ///     false
    /// }
    /// ```
    fn can_merge(&self, _other: &dyn UndoRedoCommand) -> bool {
        false
    }

    /// Merges this command with the other command.
    /// Returns true if the merge was successful.
    ///
    /// This method is called only if `can_merge` returns true.
    ///
    /// # Example
    /// ```test
    /// use common::undo_redo::UndoRedoCommand;
    ///
    /// fn merge(&mut self, other: &dyn UndoRedoCommand) -> bool {
    ///     if let Some(other_cmd) = other.as_any().downcast_ref::<Self>() {
    ///         // Merge the commands
    ///         self.value += other_cmd.value;
    ///         return true;
    ///     }
    ///     false
    /// }
    /// ```
    fn merge(&mut self, _other: &dyn UndoRedoCommand) -> bool {
        false
    }

    /// Returns the type ID of this command for type checking.
    ///
    /// This is used for downcasting in the `can_merge` and `merge` methods.
    ///
    /// # Example
    /// ```test
    /// fn as_any(&self) -> &dyn Any {
    ///     self
    /// }
    /// ```
    fn as_any(&self) -> &dyn Any;
}

/// A composite command that groups multiple commands as one.
///
/// This allows treating a sequence of commands as a single unit for undo/redo operations.
/// When a composite command is undone or redone, all its contained commands are undone
/// or redone in the appropriate order.
///
/// # Example
/// ```test
/// use common::undo_redo::CompositeCommand;
/// let mut composite = CompositeCommand::new();
/// composite.add_command(Box::new(Command1::new()));
/// composite.add_command(Box::new(Command2::new()));
/// // Now composite can be treated as a single command
/// ```
pub struct CompositeCommand {
    commands: Vec<Box<dyn UndoRedoCommand>>,
    pub stack_id: u64,
}

impl CompositeCommand {
    /// Creates a new empty composite command.
    pub fn new(stack_id: Option<u64>) -> Self {
        CompositeCommand {
            commands: Vec::new(),
            stack_id: stack_id.unwrap_or(0),
        }
    }

    /// Adds a command to this composite.
    ///
    /// Commands are executed, undone, and redone in the order they are added.
    pub fn add_command(&mut self, command: Box<dyn UndoRedoCommand>) {
        self.commands.push(command);
    }

    /// Returns true if this composite contains no commands.
    pub fn is_empty(&self) -> bool {
        self.commands.is_empty()
    }
}

impl UndoRedoCommand for CompositeCommand {
    fn undo(&mut self) -> Result<()> {
        // Undo commands in reverse order
        for command in self.commands.iter_mut().rev() {
            command.undo()?;
        }
        Ok(())
    }

    fn redo(&mut self) -> Result<()> {
        // Redo commands in original order
        for command in self.commands.iter_mut() {
            command.redo()?;
        }
        Ok(())
    }

    fn as_any(&self) -> &dyn Any {
        self
    }
}
/// Trait for commands that can be executed asynchronously with progress tracking and cancellation.
///
/// This trait extends the basic UndoRedoCommand trait with asynchronous capabilities.
/// Implementors must also implement the UndoRedoCommand trait to ensure compatibility
/// with the existing undo/redo system.
pub trait AsyncUndoRedoCommand: UndoRedoCommand {
    /// Starts the undo operation asynchronously and returns immediately.
    /// Returns Ok(()) if the operation was successfully started.
    fn start_undo(&mut self) -> Result<()>;

    /// Starts the redo operation asynchronously and returns immediately.
    /// Returns Ok(()) if the operation was successfully started.
    fn start_redo(&mut self) -> Result<()>;

    /// Checks the progress of the current operation.
    /// Returns a value between 0.0 (not started) and 1.0 (completed).
    fn check_progress(&self) -> f32;

    /// Attempts to cancel the in-progress operation.
    /// Returns Ok(()) if cancellation was successful or if no operation is in progress.
    fn cancel(&mut self) -> Result<()>;

    /// Checks if the current operation is complete.
    /// Returns true if the operation has finished successfully.
    fn is_complete(&self) -> bool;
}

#[derive(Default)]
struct StackData {
    undo_stack: Vec<Box<dyn UndoRedoCommand>>,
    redo_stack: Vec<Box<dyn UndoRedoCommand>>,
}

/// Manager for undo and redo operations.
///
/// The UndoRedoManager maintains multiple stacks of commands:
/// - Each stack has an undo stack for commands that can be undone
/// - Each stack has a redo stack for commands that have been undone and can be redone
///
/// It also supports:
/// - Grouping multiple commands as a single unit using begin_composite/end_composite
/// - Merging commands of the same type when appropriate
/// - Switching between different stacks
pub struct UndoRedoManager {
    stacks: HashMap<u64, StackData>,
    next_stack_id: u64,
    in_progress_composite: Option<CompositeCommand>,
    composite_nesting_level: usize,
    composite_stack_id: Option<u64>,
    event_hub: Option<Arc<EventHub>>,
}

impl Default for UndoRedoManager {
    fn default() -> Self {
        Self::new()
    }
}

impl UndoRedoManager {
    /// Creates a new empty UndoRedoManager with one default stack (ID 0).
    pub fn new() -> Self {
        let mut stacks = HashMap::new();
        stacks.insert(0, StackData::default());
        UndoRedoManager {
            stacks,
            next_stack_id: 1,
            in_progress_composite: None,
            composite_nesting_level: 0,
            composite_stack_id: None,
            event_hub: None,
        }
    }

    /// Inject the event hub to allow sending undo/redo related events
    pub fn set_event_hub(&mut self, event_hub: &Arc<EventHub>) {
        self.event_hub = Some(Arc::clone(event_hub));
    }

    /// Undoes the most recent command on the specified stack.
    /// If `stack_id` is None, the global stack (ID 0) is used.
    ///
    /// The undone command is moved to the redo stack.
    /// Returns Ok(()) if successful or if there are no commands to undo.
    pub fn undo(&mut self, stack_id: Option<u64>) -> Result<()> {
        let target_stack_id = stack_id.unwrap_or(0);
        let stack = self
            .stacks
            .get_mut(&target_stack_id)
            .ok_or_else(|| anyhow!("Stack with ID {} not found", target_stack_id))?;

        if let Some(mut command) = stack.undo_stack.pop() {
            if let Err(e) = command.undo() {
                log::error!("Undo failed, re-pushing command to undo stack: {e}");
                stack.undo_stack.push(command);
                return Err(e);
            }
            stack.redo_stack.push(command);
            if let Some(event_hub) = &self.event_hub {
                event_hub.send_event(Event {
                    origin: Origin::UndoRedo(UndoRedoEvent::Undone),
                    ids: Vec::<EntityId>::new(),
                    data: None,
                });
            }
        }
        Ok(())
    }

    /// Redoes the most recently undone command on the specified stack.
    /// If `stack_id` is None, the global stack (ID 0) is used.
    ///
    /// The redone command is moved back to the undo stack.
    /// Returns Ok(()) if successful or if there are no commands to redo.
    pub fn redo(&mut self, stack_id: Option<u64>) -> Result<()> {
        let target_stack_id = stack_id.unwrap_or(0);
        let stack = self
            .stacks
            .get_mut(&target_stack_id)
            .ok_or_else(|| anyhow!("Stack with ID {} not found", target_stack_id))?;

        if let Some(mut command) = stack.redo_stack.pop() {
            if let Err(e) = command.redo() {
                log::error!("Redo failed, re-pushing command to redo stack: {e}");
                stack.redo_stack.push(command);
                return Err(e);
            }
            stack.undo_stack.push(command);
            if let Some(event_hub) = &self.event_hub {
                event_hub.send_event(Event {
                    origin: Origin::UndoRedo(UndoRedoEvent::Redone),
                    ids: Vec::<EntityId>::new(),
                    data: None,
                });
            }
        }
        Ok(())
    }

    /// Begins a composite command group.
    ///
    /// All commands added between begin_composite and end_composite will be treated as a single command.
    /// This is useful for operations that logically represent a single action but require multiple
    /// commands to implement.
    ///
    /// # Example
    /// ```test
    /// let mut manager = UndoRedoManager::new();
    /// manager.begin_composite();
    /// manager.add_command(Box::new(Command1::new()));
    /// manager.add_command(Box::new(Command2::new()));
    /// manager.end_composite();
    /// // Now undo() will undo both commands as a single unit
    /// ```
    pub fn begin_composite(&mut self, stack_id: Option<u64>) -> Result<()> {
        if self.composite_stack_id.is_some() && self.composite_stack_id != stack_id {
            return Err(anyhow!(
                "Cannot begin a composite on a different stack while another composite is in progress"
            ));
        }

        // Set the target stack ID for this composite
        self.composite_stack_id = stack_id;

        // Increment the nesting level
        self.composite_nesting_level += 1;

        // If there's no composite in progress, create one
        if self.in_progress_composite.is_none() {
            self.in_progress_composite = Some(CompositeCommand::new(stack_id));
        }

        // not sure if we want to send events for composites
        if let Some(event_hub) = &self.event_hub {
            event_hub.send_event(Event {
                origin: Origin::UndoRedo(UndoRedoEvent::BeginComposite),
                ids: Vec::<EntityId>::new(),
                data: None,
            });
        }
        Ok(())
    }

    /// Ends the current composite command group and adds it to the specified undo stack.
    ///
    /// If no commands were added to the composite, nothing is added to the undo stack.
    /// If this is a nested composite, only the outermost composite is added to the undo stack.
    pub fn end_composite(&mut self) {
        // Decrement the nesting level
        if self.composite_nesting_level > 0 {
            self.composite_nesting_level -= 1;
        }

        // Only end the composite if we're at the outermost level
        if self.composite_nesting_level == 0 {
            if let Some(composite) = self.in_progress_composite.take()
                && !composite.is_empty()
            {
                let target_stack_id = self.composite_stack_id.unwrap_or(0);
                let stack = self
                    .stacks
                    .get_mut(&target_stack_id)
                    .expect("Stack must exist");
                stack.undo_stack.push(Box::new(composite));
                stack.redo_stack.clear();
            }
            // not sure if we want to send events for composites
            if let Some(event_hub) = &self.event_hub {
                event_hub.send_event(Event {
                    origin: Origin::UndoRedo(UndoRedoEvent::EndComposite),
                    ids: Vec::<EntityId>::new(),
                    data: None,
                });
            }
        }
    }

    pub fn cancel_composite(&mut self) {
        // Decrement the nesting level
        if self.composite_nesting_level > 0 {
            self.composite_nesting_level -= 1;
        }

        // Undo any sub-commands that were already executed in this composite
        if let Some(ref mut composite) = self.in_progress_composite {
            let _ = composite.undo();
        }

        self.in_progress_composite = None;
        self.composite_stack_id = None;

        // not sure if we want to send events for composites
        if let Some(event_hub) = &self.event_hub {
            event_hub.send_event(Event {
                origin: Origin::UndoRedo(UndoRedoEvent::CancelComposite),
                ids: Vec::<EntityId>::new(),
                data: None,
            });
        }
    }

    /// Adds a command to the global undo stack (ID 0).
    pub fn add_command(&mut self, command: Box<dyn UndoRedoCommand>) {
        let _ = self.add_command_to_stack(command, None);
    }

    /// Adds a command to the specified undo stack.
    /// If `stack_id` is None, the global stack (ID 0) is used.
    ///
    /// This method handles several cases:
    /// 1. If a composite command is in progress, the command is added to the composite
    /// 2. If the command can be merged with the last command on the specified undo stack, they are merged
    /// 3. Otherwise, the command is added to the specified undo stack as a new entry
    ///
    /// In all cases, the redo stack of the stack is cleared when a new command is added.
    pub fn add_command_to_stack(
        &mut self,
        command: Box<dyn UndoRedoCommand>,
        stack_id: Option<u64>,
    ) -> Result<()> {
        // If we have a composite in progress, add the command to it
        if let Some(composite) = &mut self.in_progress_composite {
            // ensure that the stack_id is the same as the composite's stack
            if composite.stack_id != stack_id.unwrap_or(0) {
                return Err(anyhow!(
                    "Cannot add command to composite with different stack ID"
                ));
            }
            composite.add_command(command);
            return Ok(());
        }

        let target_stack_id = stack_id.unwrap_or(0);
        let stack = self
            .stacks
            .get_mut(&target_stack_id)
            .ok_or_else(|| anyhow!("Stack with ID {} does not exist", target_stack_id))?;

        // Try to merge with the last command if possible
        if let Some(last_command) = stack.undo_stack.last_mut()
            && last_command.can_merge(&*command)
            && last_command.merge(&*command)
        {
            // Successfully merged, no need to add the new command
            stack.redo_stack.clear();
            return Ok(());
        }

        // If we couldn't merge, just add the command normally
        stack.undo_stack.push(command);
        stack.redo_stack.clear();
        Ok(())
    }

    /// Returns true if there are commands that can be undone on the specified stack.
    /// If `stack_id` is None, the global stack (ID 0) is used.
    pub fn can_undo(&self, stack_id: Option<u64>) -> bool {
        let target_stack_id = stack_id.unwrap_or(0);
        self.stacks
            .get(&target_stack_id)
            .map(|s| !s.undo_stack.is_empty())
            .unwrap_or(false)
    }

    /// Returns true if there are commands that can be redone on the specified stack.
    /// If `stack_id` is None, the global stack (ID 0) is used.
    pub fn can_redo(&self, stack_id: Option<u64>) -> bool {
        let target_stack_id = stack_id.unwrap_or(0);
        self.stacks
            .get(&target_stack_id)
            .map(|s| !s.redo_stack.is_empty())
            .unwrap_or(false)
    }

    /// Clears the undo and redo history for a specific stack.
    ///
    /// This method removes all commands from both the undo and redo stacks of the specified stack.
    pub fn clear_stack(&mut self, stack_id: u64) {
        if let Some(stack) = self.stacks.get_mut(&stack_id) {
            stack.undo_stack.clear();
            stack.redo_stack.clear();
        }
    }

    /// Clears all undo and redo history from all stacks.
    pub fn clear_all_stacks(&mut self) {
        for stack in self.stacks.values_mut() {
            stack.undo_stack.clear();
            stack.redo_stack.clear();
        }
        self.in_progress_composite = None;
        self.composite_nesting_level = 0;
    }

    /// Creates a new undo/redo stack and returns its ID.
    pub fn create_new_stack(&mut self) -> u64 {
        let id = self.next_stack_id;
        self.stacks.insert(id, StackData::default());
        self.next_stack_id += 1;
        id
    }

    /// Deletes an undo/redo stack by its ID.
    ///
    /// The default stack (ID 0) cannot be deleted.
    pub fn delete_stack(&mut self, stack_id: u64) -> Result<()> {
        if stack_id == 0 {
            return Err(anyhow!("Cannot delete the default stack"));
        }
        if self.stacks.remove(&stack_id).is_some() {
            Ok(())
        } else {
            Err(anyhow!("Stack with ID {} does not exist", stack_id))
        }
    }

    /// Gets the size of the undo stack for a specific stack.
    pub fn get_stack_size(&self, stack_id: u64) -> usize {
        self.stacks
            .get(&stack_id)
            .map(|s| s.undo_stack.len())
            .unwrap_or(0)
    }
}