import os
import chess.pgn

def main(pgn_path):
    os.makedirs("games", exist_ok=True)
    with open(pgn_path) as pgn_file:
        game_num = 1

        while True:
            game = chess.pgn.read_game(pgn_file)
            if game is None:
                break

            board = game.board()
            moves = []
    
            for move in game.mainline_moves():
                moves.append(move)
                board.push(move)

                if board.is_stalemate() or board.is_insufficient_material() or board.is_checkmate():
                    break
                else:
                    continue
            else:
                if not (board.is_stalemate() or board.is_insufficient_material() or board.is_checkmate()):
                    continue

            new_game = chess.pgn.Game(headers={'Result': game.headers['Result']})

            node = new_game
            for move in moves:
                node = node.add_variation(move)

            out_path = os.path.join("games", f"{game_num:04d}.pgn")
            with open(out_path, "w") as out:
                out.write(str(new_game))
            
            game_num += 1

if __name__ == "__main__":
    main("twic1621.pgn")