import pygame
import sys

# 初始化Pygame
pygame.init()

# 定义常量
SCREEN_WIDTH = 600
SCREEN_HEIGHT = 600
ROW_COUNT = 15
COL_COUNT = 15
SQUARE_SIZE = SCREEN_WIDTH // ROW_COUNT

# 定义颜色
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
RED = (255, 0, 0)

# 创建屏幕
screen = pygame.display.set_mode((SCREEN_WIDTH, SCREEN_HEIGHT))
pygame.display.set_caption('五子棋')


# 绘制网格
def draw_grid():
    for row in range(ROW_COUNT):
        pygame.draw.line(screen, BLACK, (0, row * SQUARE_SIZE), (SCREEN_WIDTH, row * SQUARE_SIZE))
    for col in range(COL_COUNT):
        pygame.draw.line(screen, BLACK, (col * SQUARE_SIZE, 0), (col * SQUARE_SIZE, SCREEN_HEIGHT))

# 绘制棋子
def draw_piece(row, col, color):
    pygame.draw.circle(screen, color, (col * SQUARE_SIZE + SQUARE_SIZE // 2, row * SQUARE_SIZE + SQUARE_SIZE // 2),
                       SQUARE_SIZE // 2 - 5)

# # 绘制悔棋和重新开始按钮
# def draw_buttons():
#     undo_text = font.render('悔棋 (u)', True, BLACK)
#     restart_text = font.render('重新开始 (r)', True, BLACK)
#     screen.blit(undo_text, (10, SCREEN_HEIGHT - 60))
#     screen.blit(restart_text, (10, SCREEN_HEIGHT - 30))
# 检查胜利条件
def check_win(board, row, col, color):
    # 水平方向
    count = 1
    for i in range(1, 5):
        if col + i < COL_COUNT and board[row][col + i] == color:
            count += 1
        else:
            break
    for i in range(1, 5):
        if col - i >= 0 and board[row][col - i] == color:
            count += 1
        else:
            break
    if count >= 5:
        return True

    # 垂直方向
    count = 1
    for i in range(1, 5):
        if row + i < ROW_COUNT and board[row + i][col] == color:
            count += 1
        else:
            break
    for i in range(1, 5):
        if row - i >= 0 and board[row - i][col] == color:
            count += 1
        else:
            break
    if count >= 5:
        return True

    # 左上到右下
    count = 1
    for i in range(1, 5):
        if row + i < ROW_COUNT and col + i < COL_COUNT and board[row + i][col + i] == color:
            count += 1
        else:
            break
    for i in range(1, 5):
        if row - i >= 0 and col - i >= 0 and board[row - i][col - i] == color:
            count += 1
        else:
            break
    if count >= 5:
        return True

    # 右上到左下
    count = 1
    for i in range(1, 5):
        if row - i >= 0 and col + i < COL_COUNT and board[row - i][col + i] == color:
            count += 1
        else:
            break
    for i in range(1, 5):
        if row + i < ROW_COUNT and col - i >= 0 and board[row + i][col - i] == color:
            count += 1
        else:
            break
    if count >= 5:
        return True

    return False


# 初始化棋盘
def init_board():
    board = [[0 for _ in range(COL_COUNT)] for _ in range(ROW_COUNT)]
    return board


# 主函数
def main():
    board = init_board()
    turn = 0  # 0: 黑子, 1: 白子
    game_over = False
    undo_board = []

    font = pygame.font.SysFont(None, 40)

    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                sys.exit()

                # if not game_over and event.type == pygame.MOUSEBUTTONDOWN:
                #     mouseX, mouseY = event.pos
                #     if 10 <= mouseX <= 100 and SCREEN_HEIGHT - 60 <= mouseY <= SCREEN_HEIGHT - 30:  # 点击了悔棋按钮
                #         if len(undo_board) > 1:
                #             board = undo_board.pop()
                #             turn = 1 - turn
                #     elif 10 <= mouseX <= 170 and SCREEN_HEIGHT - 30 <= mouseY <= SCREEN_HEIGHT:  # 点击了重新开始按钮
                #         board = init_board()
                #         turn = 0
                #         game_over = False
                #         undo_board.clear()
                #
                #     else:
                #         clicked_row = mouseY // SQUARE_SIZE
                #         clicked_col = mouseX // SQUARE_SIZE
                #         if board[clicked_row][clicked_col] == 0:
                #             undo_board.append([list(row) for row in board])
                #             if turn == 0:
                #                 board[clicked_row][clicked_col] = 1
                #                 if check_win(board, clicked_row, clicked_col, 1):
                #                     game_over = True
                #             else:
                #                 board[clicked_row][clicked_col] = 2
                #                 if check_win(board, clicked_row, clicked_col, 2):
                #                     game_over = True
                #             turn = 1 - turn

            if not game_over and event.type == pygame.MOUSEBUTTONDOWN:
                mouseX, mouseY = event.pos
                clicked_row = mouseY // SQUARE_SIZE
                clicked_col = mouseX // SQUARE_SIZE
                if board[clicked_row][clicked_col] == 0:
                    undo_board.append([list(row) for row in board])
                    if turn == 0:
                        board[clicked_row][clicked_col] = 1
                        if check_win(board, clicked_row, clicked_col, 1):
                            game_over = True
                    else:
                        board[clicked_row][clicked_col] = 2
                        if check_win(board, clicked_row, clicked_col, 2):
                            game_over = True
                    turn = 1 - turn

            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_u:  # 悔棋
                    if len(undo_board) > 1:
                        board = undo_board.pop()
                        turn = 1 - turn
                elif event.key == pygame.K_r:  # 重新开始
                    board = init_board()
                    turn = 0
                    game_over = False
                    undo_board.clear()

        screen.fill(WHITE)
        draw_grid()
        # draw_buttons()

        for row in range(ROW_COUNT):
            for col in range(COL_COUNT):
                if board[row][col] == 1:
                    draw_piece(row, col, BLACK)
                elif board[row][col] == 2:
                    draw_piece(row, col, RED)

        current_player_text = font.render(f'当前玩家: {"黑子" if turn == 0 else "白子"}', True, BLACK)
        screen.blit(current_player_text, (10, 10))

        if game_over:
            winner_text = font.render(f'{"黑子" if turn == 1 else "白子"}胜利!', True, RED)
            screen.blit(winner_text, (200, SCREEN_HEIGHT // 2 - 20))

        pygame.display.update()


# 运行游戏
if __name__ == '__main__':
    main()