From 91678a42b2012b80e3716ca245ca152030e00588 Mon Sep 17 00:00:00 2001
From: Persson-dev <sim16.prib@gmail.com>
Date: Mon, 14 Apr 2025 11:02:35 +0200
Subject: [PATCH] fix undo draw check

---
 .../controller/commands/MoveCommand.java      |  2 ++
 .../controller/commands/NewGameCommand.java   |  2 +-
 .../controller/commands/PromoteCommand.java   |  2 ++
 app/src/main/java/chess/model/Game.java       | 32 ++++++++++++++++---
 4 files changed, 33 insertions(+), 5 deletions(-)

diff --git a/app/src/main/java/chess/controller/commands/MoveCommand.java b/app/src/main/java/chess/controller/commands/MoveCommand.java
index 5a1b159..bd1b8ba 100644
--- a/app/src/main/java/chess/controller/commands/MoveCommand.java
+++ b/app/src/main/java/chess/controller/commands/MoveCommand.java
@@ -34,6 +34,7 @@ public class MoveCommand extends PlayerCommand {
 
 			case Moved:
 				outputSystem.onMove(this.move);
+				game.saveTraitPiecesPos();
 				return result;
 
 			case ActionNeeded:
@@ -83,6 +84,7 @@ public class MoveCommand extends PlayerCommand {
 	protected CommandResult undoImpl(Game game, GameListener outputSystem) {
 		final ChessBoard board = game.getBoard();
 
+		game.undoTraitPiecesPos();
 		board.undoMove(move, deadPiece);
 		return CommandResult.Moved;
 	}
diff --git a/app/src/main/java/chess/controller/commands/NewGameCommand.java b/app/src/main/java/chess/controller/commands/NewGameCommand.java
index 51bb373..07a88bf 100644
--- a/app/src/main/java/chess/controller/commands/NewGameCommand.java
+++ b/app/src/main/java/chess/controller/commands/NewGameCommand.java
@@ -48,7 +48,7 @@ public class NewGameCommand extends Command {
 		board.pieceComes(new Queen(Color.White), new Coordinate(3, Coordinate.VALUE_MAX - 1));
 		board.pieceComes(new King(Color.White), new Coordinate(4, Coordinate.VALUE_MAX - 1));
 
-		game.resetPlayerTurn();
+		game.reset();
 
 		outputSystem.onGameStart();
 		outputSystem.onPlayerTurn(game.getPlayerTurn());
diff --git a/app/src/main/java/chess/controller/commands/PromoteCommand.java b/app/src/main/java/chess/controller/commands/PromoteCommand.java
index 4df4eb9..ebf9136 100644
--- a/app/src/main/java/chess/controller/commands/PromoteCommand.java
+++ b/app/src/main/java/chess/controller/commands/PromoteCommand.java
@@ -85,6 +85,8 @@ public class PromoteCommand extends PlayerCommand {
 
 		assert promoted != null;
 
+		game.undoTraitPiecesPos();
+
 		board.pieceComes(this.oldPawn, this.pieceCoords);
 
 		game.getLastAction().undo(game, outputSystem);
diff --git a/app/src/main/java/chess/model/Game.java b/app/src/main/java/chess/model/Game.java
index de842f5..6b976c8 100644
--- a/app/src/main/java/chess/model/Game.java
+++ b/app/src/main/java/chess/model/Game.java
@@ -23,8 +23,6 @@ public class Game {
 		this.board = board;
 		this.movesHistory = new Stack<>();
 		this.traitsPos = new HashMap<>();
-		this.traitsPos.put(Color.Black, new HashMap<>());
-		this.traitsPos.put(Color.White, new HashMap<>());
 	}
 
 	public ChessBoard getBoard() {
@@ -35,6 +33,18 @@ public class Game {
 		return playerTurn;
 	}
 
+	public void reset() {
+		resetPlayerTurn();
+		this.traitsPos.clear();
+		initPlayer(Color.Black);
+		initPlayer(Color.White);
+	}
+
+	private void initPlayer(Color color) {
+		this.traitsPos.put(color, new HashMap<>());
+		saveTraitPiecesPos(color);
+	}
+
 	public void resetPlayerTurn() {
 		this.playerTurn = Color.White;
 	}
@@ -50,6 +60,10 @@ public class Game {
 		this.traitsPos.get(color).put(piecesHash, count == null ? 1 : count + 1);
 	}
 
+	public void saveTraitPiecesPos() {
+		saveTraitPiecesPos(playerTurn);
+	}
+
 	/**
 	 * 
 	 * @param player the trait to check
@@ -73,14 +87,13 @@ public class Game {
 	 * @return true if a draw should occur
 	 */
 	public void switchPlayerTurn() {
-		saveTraitPiecesPos(this.playerTurn);
 		playerTurn = Color.getEnemy(playerTurn);
 	}
 
 	public GameStatus checkGameStatus() {
 		final Color enemy = Color.getEnemy(getPlayerTurn());
 
-		if(checkDraw())
+		if (checkDraw())
 			return GameStatus.Draw;
 
 		if (this.board.isKingInCheck(enemy))
@@ -115,4 +128,15 @@ public class Game {
 		}
 	}
 
+	private void undoTraitPiecesPos(Color color) {
+		int piecesHash = this.board.hashPlayerPieces(color);
+		Integer count = this.traitsPos.get(color).get(piecesHash);
+		if (count != null)
+			this.traitsPos.get(color).put(piecesHash, count - 1);
+	}
+
+	public void undoTraitPiecesPos() {
+		undoTraitPiecesPos(Color.getEnemy(playerTurn));
+	}
+
 }