1 files changed, 49 insertions, 5 deletions

diff --git a/src/main/java/design/model/holeplay/HolePlayContext.java b/src/main/java/design/model/holeplay/HolePlayContext.java
index 16201b6..71fba32 100644
--- a/src/main/java/design/model/holeplay/HolePlayContext.java
+++ b/src/main/java/design/model/holeplay/HolePlayContext.java
@@ -113,15 +113,59 @@ public class HolePlayContext implements Originator {
         this.strokes = s;
     }
 
+    // Captures current State, current play (hole# + swings), and stroke counter
+    private static class HolePlayMemento implements Memento {
+        private final String stateName;
+        private final Play currentPlayCopy;
+        private final int strokesCopy;
+
+        HolePlayMemento(HolePlayContext ctx) {
+            this.stateName = ctx.state.name();
+
+            // Copy currentPlay if it exists
+            Play src = ctx.currentPlay;
+            if (src != null) {
+                var copySwings = new ArrayList<Swing>();
+                for (Swing s : src.getSwings()) {
+                    copySwings.add(new Swing(s.getDistance(), s.getClubUsed()));
+                }
+                this.currentPlayCopy = new Play(src.getHoleNumber(), copySwings);
+            } else {
+                this.currentPlayCopy = null;
+            }
+
+            this.strokesCopy = ctx.strokes;
+        }
+    }
+
     @Override
     public Memento createMemento() {
-        // TODO Auto-generated method stub
-        throw new UnsupportedOperationException("Unimplemented method 'createMemento'");
+        return new HolePlayMemento(this);
     }
 
     @Override
-    public void restore(Memento memento) {
-        // TODO Auto-generated method stub
-        throw new UnsupportedOperationException("Unimplemented method 'restore'");
+    public void restore(Memento m) {
+        HolePlayMemento hm = (HolePlayMemento) m;
+
+        // restore strokes
+        this.strokes = hm.strokesCopy;
+
+        // restore currentPlay
+        if (hm.currentPlayCopy == null) {
+            this.currentPlay = null;
+        } else {
+            var copySwings = new ArrayList<Swing>();
+            for (Swing s : hm.currentPlayCopy.getSwings()) {
+                copySwings.add(new Swing(s.getDistance(), s.getClubUsed()));
+            }
+            this.currentPlay = new Play(hm.currentPlayCopy.getHoleNumber(), copySwings);
+        }
+
+        // restore state based on name
+        switch (hm.stateName) {
+            case "SetupState" -> this.state = new SetupState();
+            case "PlayState" -> this.state = new PlayState();
+            case "HoleCompleteState" -> this.state = new HoleCompleteState();
+        }
     }
 }
\ No newline at end of file