""" Undo/Redo manager for OCCSketch using snapshot-based approach. python_solvespace has no per-entity delete API — the solver is rebuilt from scratch after every modification. This makes snapshot-based undo natural: we capture the complete sketch state (entities, geometry, constraints) and restore it by rebuilding the solver from the snapshot data. Each snapshot is ~10-50 KB depending on sketch complexity, and we cap the stack at a configurable depth (default 50). """ from __future__ import annotations import copy import logging from dataclasses import dataclass, field from typing import Any, Dict, List, Optional, Set, Tuple logger = logging.getLogger(__name__) @dataclass class SketchSnapshot: """Immutable snapshot of sketch state for undo/redo. Captures everything needed to fully reconstruct an OCCSketch: entities, points, lines, circles, arcs, counter, constraint log, and the special entity id sets (centerlines, external/underlay). """ # Entity registry: id → (entity_type, geometry, is_construction, is_external, constraints_list) entities: Dict[int, Tuple[str, Any, bool, bool, List[str]]] = field(default_factory=dict) # Geometry sub-indices points: Dict[int, Tuple[float, float]] = field(default_factory=dict) lines: Dict[int, Tuple[int, int]] = field(default_factory=dict) circles: Dict[int, Tuple[int, float]] = field(default_factory=dict) arcs: Dict[int, Dict[str, Any]] = field(default_factory=dict) # Counters and flags entity_counter: int = 0 constraint_count: int = 0 first_point_id: Optional[int] = None # Constraint replay log constraint_log: List[Dict[str, Any]] = field(default_factory=list) # Special entity sets centerline_ids: Set[int] = field(default_factory=set) external_entity_ids: Set[int] = field(default_factory=set) # Workplane (so undo doesn't lose the placement plane) wp_origin: Tuple[float, float, float] = (0.0, 0.0, 0.0) wp_normal: Tuple[float, float, float] = (0.0, 0.0, 1.0) wp_x_dir: Tuple[float, float, float] = (1.0, 0.0, 0.0) wp_y_dir: Tuple[float, float, float] = (0.0, 1.0, 0.0) class SketchUndoManager: """Manages undo/redo stacks of SketchSnapshot for an OCCSketch. Usage:: undo_mgr = SketchUndoManager(sketch) # Before any modification: undo_mgr.save_state() # ... perform modification ... # Ctrl+Z: undo_mgr.undo() # Ctrl+Y / Ctrl+Shift+Z: undo_mgr.redo() """ def __init__(self, sketch: Any, max_stack_size: int = 50) -> None: from fluency.geometry_occ.sketch import OCCSketch self._sketch: OCCSketch = sketch self._max_stack_size = max_stack_size self._undo_stack: List[SketchSnapshot] = [] self._redo_stack: List[SketchSnapshot] = [] # ─── Public API ──────────────────────────────────────────────────────── @property def can_undo(self) -> bool: """True if there is a state to undo to.""" return len(self._undo_stack) > 0 @property def can_redo(self) -> bool: """True if there is a state to redo to.""" return len(self._redo_stack) > 0 @property def undo_depth(self) -> int: """Number of undo levels available.""" return len(self._undo_stack) @property def redo_depth(self) -> int: """Number of redo levels available.""" return len(self._redo_stack) def save_state(self) -> None: """Capture the current sketch state and push it onto the undo stack. Call this **before** any modifying operation (draw, delete, move, constraint add/remove, construction toggle, etc.). The redo stack is cleared whenever a new state is saved (i.e. when the user makes a new change after undoing). """ snapshot = self._capture() self._undo_stack.append(snapshot) # Cap the stack size. if len(self._undo_stack) > self._max_stack_size: self._undo_stack.pop(0) # New mutation invalidates the redo history. self._redo_stack.clear() logger.debug( f"save_state: undo_depth={len(self._undo_stack)} " f"entities={len(snapshot.entities)}" ) def undo(self) -> bool: """Restore the previous sketch state. Returns True if a state was restored, False if the undo stack is empty. """ if not self._undo_stack: logger.debug("undo: stack empty") return False # Save current state to redo stack before restoring. current_snapshot = self._capture() self._redo_stack.append(current_snapshot) # Pop and restore. snapshot = self._undo_stack.pop() self._restore(snapshot) logger.debug( f"undo: restored state with {len(snapshot.entities)} entities, " f"undo_depth={len(self._undo_stack)} redo_depth={len(self._redo_stack)}" ) return True def redo(self) -> bool: """Re-apply the most recently undone state. Returns True if a state was restored, False if the redo stack is empty. """ if not self._redo_stack: logger.debug("redo: stack empty") return False # Save current state to undo stack before restoring. current_snapshot = self._capture() self._undo_stack.append(current_snapshot) # Pop and restore. snapshot = self._redo_stack.pop() self._restore(snapshot) logger.debug( f"redo: restored state with {len(snapshot.entities)} entities, " f"undo_depth={len(self._undo_stack)} redo_depth={len(self._redo_stack)}" ) return True def clear(self) -> None: """Clear both stacks (e.g. when loading a new sketch).""" self._undo_stack.clear() self._redo_stack.clear() logger.debug("undo stacks cleared") # ─── Snapshot Capture ────────────────────────────────────────────────── def _capture(self) -> SketchSnapshot: """Capture the current sketch state into a snapshot.""" sketch = self._sketch # Capture entities: id → (type, geometry, is_construction, is_external, constraints) entities: Dict[int, Tuple[str, Any, bool, bool, List[str]]] = {} for eid, ent in sketch._entities.items(): entities[eid] = ( ent.entity_type, ent.geometry, ent.is_construction, ent.is_external, list(ent.constraints), # copy the constraints list ) # Deep copy the mutable dicts (points coords are tuples, so shallow is fine, # but arcs contain dicts so we deep-copy those). points = dict(sketch._points) lines = dict(sketch._lines) circles = dict(sketch._circles) arcs = {k: copy.deepcopy(v) for k, v in sketch._arcs.items()} # Constraint log entries contain tuples and sets — need careful copy. constraint_log = [] for entry in sketch._constraint_log: copied = { "type": entry["type"], "ids": tuple(entry["ids"]), "params": tuple(entry["params"]) if entry.get("params") else (), "labels": set(entry["labels"]) if entry.get("labels") else set(), } constraint_log.append(copied) return SketchSnapshot( entities=entities, points=points, lines=lines, circles=circles, arcs=arcs, entity_counter=sketch._entity_counter, constraint_count=sketch._constraint_count, first_point_id=sketch._first_point_id, constraint_log=constraint_log, centerline_ids=set(sketch._centerline_ids), external_entity_ids=set(sketch._external_entity_ids), wp_origin=sketch._wp_origin, wp_normal=sketch._wp_normal, wp_x_dir=sketch._wp_x_dir, wp_y_dir=sketch._wp_y_dir, ) # ─── Snapshot Restore ────────────────────────────────────────────────── def _restore(self, snapshot: SketchSnapshot) -> None: """Restore the sketch to a previously captured snapshot state.""" from fluency.geometry_occ.sketch import OCCSketch, OCCSketchEntity sketch = self._sketch # Clear the current solver and rebuild from scratch. sketch._solver = sketch._solver.__class__() # SolverSystem() sketch._wp = sketch._solver.create_2d_base() sketch._first_point_id = None # Restore counters and flags. sketch._entity_counter = snapshot.entity_counter sketch._constraint_count = snapshot.constraint_count sketch._centerline_ids = set(snapshot.centerline_ids) sketch._external_entity_ids = set(snapshot.external_entity_ids) # Restore workplane. sketch._wp_origin = snapshot.wp_origin sketch._wp_normal = snapshot.wp_normal sketch._wp_x_dir = snapshot.wp_x_dir sketch._wp_y_dir = snapshot.wp_y_dir # Rebuild entity objects from the snapshot. sketch._entities.clear() sketch._points.clear() sketch._lines.clear() sketch._circles.clear() sketch._arcs.clear() # First pass: re-add all points to the solver. for eid in sorted(snapshot.entities.keys()): etype, geometry, is_constr, is_ext, constraints = snapshot.entities[eid] if etype == "point" and eid in snapshot.points: x, y = snapshot.points[eid] solver_handle = sketch._solver.add_point_2d(x, y, sketch._wp) ent = OCCSketchEntity( entity_id=eid, entity_type="point", geometry=(x, y), handle=solver_handle, ) ent.is_construction = is_constr ent.is_external = is_ext ent.constraints = list(constraints) sketch._entities[eid] = ent sketch._points[eid] = (x, y) # Anchor the first point (or first external point) for solver stability. if sketch._first_point_id is None and not is_ext: sketch._first_point_id = eid sketch._solver.dragged(solver_handle, sketch._wp) elif sketch._first_point_id is None and is_ext: sketch._first_point_id = eid sketch._solver.dragged(solver_handle, sketch._wp) # Second pass: re-add all lines. for lid in sorted(snapshot.lines.keys()): sid, eid2 = snapshot.lines[lid] s_ent = sketch._entities.get(sid) e_ent = sketch._entities.get(eid2) if s_ent is None or e_ent is None or s_ent.handle is None or e_ent.handle is None: continue solver_handle = sketch._solver.add_line_2d(s_ent.handle, e_ent.handle, sketch._wp) etype, geometry, is_constr, is_ext, constraints = snapshot.entities[lid] ent = OCCSketchEntity( entity_id=lid, entity_type="line", geometry=geometry, handle=solver_handle, ) ent.is_construction = is_constr ent.is_external = is_ext ent.constraints = list(constraints) sketch._entities[lid] = ent sketch._lines[lid] = (sid, eid2) # Restore circles (not in solver, just tracked). for cid, (center_id, radius) in snapshot.circles.items(): if cid in snapshot.entities: etype, geometry, is_constr, is_ext, constraints = snapshot.entities[cid] center_ent = sketch._entities.get(center_id) ent = OCCSketchEntity( entity_id=cid, entity_type="circle", geometry=geometry, handle=None, ) ent.is_construction = is_constr ent.is_external = is_ext ent.constraints = list(constraints) sketch._entities[cid] = ent sketch._circles[cid] = (center_id, radius) # Restore arcs (not in solver, just tracked). for aid, arc_data in snapshot.arcs.items(): if aid in snapshot.entities: etype, geometry, is_constr, is_ext, constraints = snapshot.entities[aid] ent = OCCSketchEntity( entity_id=aid, entity_type="arc", geometry=geometry, handle=None, ) ent.is_construction = is_constr ent.is_external = is_ext ent.constraints = list(constraints) sketch._entities[aid] = ent sketch._arcs[aid] = copy.deepcopy(arc_data) # Rebuild the constraint log (entries were deep-copied on capture). sketch._constraint_log = [] for entry in snapshot.constraint_log: sketch._constraint_log.append({ "type": entry["type"], "ids": tuple(entry["ids"]), "params": tuple(entry["params"]) if entry.get("params") else (), "labels": set(entry["labels"]) if entry.get("labels") else set(), }) # Re-apply all constraints to the solver. for entry in sketch._constraint_log: sketch._apply_constraint_log(entry) # Solve to update geometry positions. sketch.solve() logger.debug( f"Restored snapshot: {len(sketch._entities)} entities, " f"{len(sketch._constraint_log)} constraints" )