from typing import Dict, Any, List
from copy import deepcopy

class WorkflowState:
    """Manages state and history for workflow execution"""
    
    def __init__(self):
        self.context: Dict[str, Any] = {}
        self.history: List[Dict[str, Any]] = []
        
    def initialize(self, initial_context: Dict[str, Any]):
        """Initialize workflow state with initial context"""
        self.context = deepcopy(initial_context)
        self.history = []
        
    def update(self, updates: Dict[str, Any]):
        """Update workflow state with new data"""
        self.context.update(deepcopy(updates))
        self.history.append(deepcopy(updates))
        
    def get_context(self) -> Dict[str, Any]:
        """Get current workflow context"""
        return deepcopy(self.context)
    
    def get_history(self) -> List[Dict[str, Any]]:
        """Get workflow execution history"""
        return deepcopy(self.history)
    
    def rollback(self, steps: int = 1):
        """Rollback workflow state by specified number of steps"""
        if steps > len(self.history):
            steps = len(self.history)
            
        # Remove recent history entries
        for _ in range(steps):
            self.history.pop()
            
        # Reconstruct context from remaining history
        self.context = {}
        for entry in self.history:
            self.context.update(entry)
