4 files changed, 24 insertions, 44 deletions

diff --git a/src/board.cpp b/src/board.cpp
index 4a9531a..eb42769 100644
--- a/src/board.cpp
+++ b/src/board.cpp
@@ -48,7 +48,6 @@ UndoMove MakeMove(Move move, Game *g) {
   undo.OldblackCastleQueen = g->blackCastleQueen;
 
   // play move
-
   Piece piece = g->pieces[PositionToIndex(move.From)];
   Piece piece2 = g->pieces[PositionToIndex(move.To)];
 
diff --git a/src/bot.cpp b/src/bot.cpp
index f1f1fd6..b3e75d6 100644
--- a/src/bot.cpp
+++ b/src/bot.cpp
@@ -221,13 +221,19 @@ int PSTIndex(int square, bool white) { return white ? square : (56 ^ square); }
 float EvaluateBoardForWhite(Game *g, int depth) {
   float score = 0;
 
-  if (g->state == WHITE_WON) {
+  GameState state = GetNewGameState(g);
+  switch (state) {
+  case WHITE_WON:
     return MATE + depth;
-  }
-  if (g->state == BLACK_WON) {
+    break;
+  case BLACK_WON:
     return -MATE - depth;
-  }
-  if (g->state == STALEMATE || g->state == DRAW) {
+    break;
+  case TURN:
+    break;
+  case STALEMATE:
+    return 0;
+  case DRAW:
     return 0;
   }
 
diff --git a/src/moves.cpp b/src/moves.cpp
index 2663ea0..f741b7f 100644
--- a/src/moves.cpp
+++ b/src/moves.cpp
@@ -164,6 +164,17 @@ std::vector<Move> GetLegalMoves(Game *g) {
     }
   }
 
+  return legalMoves;
+}
+GameState GetNewGameState(Game *g) {
+  auto legalMoves = GetLegalMoves(g);
+
+  // make sure we dont override game ending states
+  if (g->state == DRAW || g->state == STALEMATE || g->state == WHITE_WON ||
+      g->state == BLACK_WON) {
+    return g->state;
+  }
+
   if (legalMoves.empty()) {
     Position kingPosition = FindKing(g, g->turn);
 
@@ -175,44 +186,7 @@ std::vector<Move> GetLegalMoves(Game *g) {
       g->state = STALEMATE;
     }
   }
-
-  for (Move move : moves) {
-    bool isLegal = true;
-    UndoMove undo = MakeMove(move, g);
-    auto opponentResponse = GetPseudoLegalMoves(g);
-    for (Move move : opponentResponse) {
-      if (g->pieces[PositionToIndex(move.To)].type == KING) {
-        isLegal = false;
-      }
-    }
-    UnMakeMove(undo, g);
-    if (isLegal) {
-      legalMoves.push_back(move);
-    };
-  }
-  if (legalMoves.size() == 0) {
-    bool isCheck = false;
-
-    Game testBoard = *g;
-    testBoard.turn = !testBoard.turn;
-    auto opponentResponse = GetPseudoLegalMoves(&testBoard);
-    for (Move move : opponentResponse) {
-      if (testBoard.pieces[PositionToIndex(move.To)].type == KING) {
-        isCheck = true;
-      }
-    }
-
-    if (isCheck) {
-      if (g->turn) {
-        g->state = BLACK_WON;
-      } else {
-        g->state = WHITE_WON;
-      };
-    } else {
-      g->state = STALEMATE;
-    }
-  }
-  return legalMoves;
+  return TURN;
 }
 
 bool IsSquareAttacked(Game *board, Position square, bool white) {
diff --git a/src/moves.hpp b/src/moves.hpp
index 20ef0ba..3b84afe 100644
--- a/src/moves.hpp
+++ b/src/moves.hpp
@@ -14,5 +14,6 @@ void GenerateKingMoves(Game *g, int from, std::vector<Move> &moves);
 void GenerateKnightMoves(Game *g, int from, std::vector<Move> &moves);
 void GeneratePawnMoves(Game *g, int from, std::vector<Move> &moves);
 bool IsSquareAttacked(Game *g, Position square, bool byWhite);
+GameState GetNewGameState(Game *g);
 
 #endif /* SRC_MOVES_H_ */