1 files changed, 12 insertions, 1 deletions
diff --git a/src/main/java/design/controller/userinput/menus/HolePlayMenu.java b/src/main/java/design/controller/userinput/menus/HolePlayMenu.java
index e0d1572..2a7fc5d 100644
--- a/src/main/java/design/controller/userinput/menus/HolePlayMenu.java
+++ b/src/main/java/design/controller/userinput/menus/HolePlayMenu.java
@@ -3,12 +3,14 @@ package design.controller.userinput.menus;
 import design.controller.userinput.Menu;
 import design.controller.userinput.MenuOption;
 import design.model.Course;
+import design.model.Golfer;
 import design.model.Hole;
 import design.runtime.Session;
 import design.model.Club;
 import design.model.Round;
 import design.model.holeplay.HolePlayContext;
 import design.persistence.PersonalDatabase;
+import design.model.undo.UndoManager;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -18,10 +20,11 @@ public class HolePlayMenu extends Menu {
 
     private final Round round;
     private final HolePlayContext ctx;
+    private final Golfer golfer = Session.getCurrentGolfer();
 
     public HolePlayMenu() {
         this.round = createRound();
-        this.ctx = new HolePlayContext(Session.getCurrentGolfer(), round, PersonalDatabase.instance());
+        this.ctx = new HolePlayContext(this.golfer, round, PersonalDatabase.instance());
     }
 
     @Override
@@ -60,6 +63,8 @@ public class HolePlayMenu extends Menu {
                     dist = Integer.parseInt(ds);
                 }
 
+                UndoManager.instance().capture(ctx, "Shot with " + club.getNickname() + " for " + dist + " yds");
+
                 // Records shot
                 ctx.recordShot(club, dist);
                 System.out.println("Shot recorded: " + dist + " yds with " + club.getNickname());
@@ -71,6 +76,9 @@ public class HolePlayMenu extends Menu {
             opts.add(new MenuOption("hole out", () -> {
                 // Precedes to next hole
                 int prev = round.getCurrentHole().getNumber();
+
+                UndoManager.instance().capture(golfer, "Hole out on " + prev + " at " + round.getCourse().getName());
+
                 ctx.holeOut();
                 System.out.println("Holed out on " + prev + ". Next: " + round.getCurrentHole().getNumber());
                 this.present();
@@ -79,6 +87,7 @@ public class HolePlayMenu extends Menu {
 
         // End round (always shown)
         opts.add(new MenuOption("end round", () -> {
+            UndoManager.instance().capture(golfer, "End round on " + round.getCourse().getName());
             ctx.endRoundNow();
             System.out.println("Round ended.");
             new MainMenu().present();
@@ -118,6 +127,8 @@ public class HolePlayMenu extends Menu {
         // Starts round and sends user to HolePlayMenu
         Hole startHole = course.getHoles().get(startHoleNum - 1);
         Round r = new Round(course, Session.getDateTime(), startHole);
+        // Undo puts golfer back to before round existed.
+        UndoManager.instance().capture(golfer, "Start round on " + course.getName() + " (hole " + startHoleNum + ")");
         golfer.addRound(r);
         return r;
     }