import pygame
import sys
import math

# ---------------------- 初始化 pygame ----------------------
pygame.init()
WIDTH, HEIGHT = 600, 650
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("井字棋 - AI对手")

# 颜色定义
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
BLUE = (0, 120, 255)
RED = (255, 50, 50)
GRAY = (180, 180, 180)
GREEN = (30, 180, 80)

# 字体（兼容中文）
try:
    font = pygame.font.Font("simhei.ttf", 40)
    small_font = pygame.font.Font("simhei.ttf", 28)
except:
    font = pygame.font.SysFont("microsoftyahei", 40)
    small_font = pygame.font.SysFont("microsoftyahei", 28)

# 棋盘设置
CELL_SIZE = 200
board = [[None, None, None],
         [None, None, None],
         [None, None, None]]
player = "X"   # 玩家 X
ai = "O"       # AI O
game_over = False
winner = None

# ---------------------- 极小极大AI算法 ----------------------
def check_win(b, mark):
    # 行
    for row in b:
        if row[0] == row[1] == row[2] == mark:
            return True
    # 列
    for col in range(3):
        if b[0][col] == b[1][col] == b[2][col] == mark:
            return True
    # 对角线
    if b[0][0] == b[1][1] == b[2][2] == mark:
        return True
    if b[0][2] == b[1][1] == b[2][0] == mark:
        return True
    return False

def is_board_full(b):
    for row in b:
        if None in row:
            return False
    return True

def minimax(b, depth, is_max):
    if check_win(b, ai):
        return 10 - depth
    if check_win(b, player):
        return depth - 10
    if is_board_full(b):
        return 0

    if is_max:
        best_score = -math.inf
        for i in range(3):
            for j in range(3):
                if b[i][j] is None:
                    b[i][j] = ai
                    score = minimax(b, depth + 1, False)
                    b[i][j] = None
                    best_score = max(best_score, score)
        return best_score
    else:
        best_score = math.inf
        for i in range(3):
            for j in range(3):
                if b[i][j] is None:
                    b[i][j] = player
                    score = minimax(b, depth + 1, True)
                    b[i][j] = None
                    best_score = min(best_score, score)
        return best_score

def get_ai_move():
    best_score = -math.inf
    best_move = None
    for i in range(3):
        for j in range(3):
            if board[i][j] is None:
                board[i][j] = ai
                score = minimax(board, 0, False)
                board[i][j] = None
                if score > best_score:
                    best_score = score
                    best_move = (i, j)
    return best_move

# ---------------------- 绘制函数 ----------------------
def draw_board():
    screen.fill(WHITE)
    # 画网格线
    for i in range(1, 3):
        pygame.draw.line(screen, BLACK, (i*CELL_SIZE, 0), (i*CELL_SIZE, HEIGHT-50), 3)
        pygame.draw.line(screen, BLACK, (0, i*CELL_SIZE), (WIDTH, i*CELL_SIZE), 3)

    # 画棋子
    for i in range(3):
        for j in range(3):
            cx = j * CELL_SIZE + CELL_SIZE//2
            cy = i * CELL_SIZE + CELL_SIZE//2
            if board[i][j] == "X":
                text = font.render("X", True, BLUE)
                rect = text.get_rect(center=(cx, cy))
                screen.blit(text, rect)
            elif board[i][j] == "O":
                text = font.render("O", True, RED)
                rect = text.get_rect(center=(cx, cy))
                screen.blit(text, rect)

    # 底部状态文字
    status_text = ""
    if game_over:
        if winner == player:
            status_text = "你赢了！"
        elif winner == ai:
            status_text = "AI获胜！"
        else:
            status_text = "平局！"
    else:
        status_text = "你的回合(X)"

    status_surf = small_font.render(status_text, True, BLACK)
    screen.blit(status_surf, (20, 610))

    # 重新开始按钮
    btn_rect = pygame.Rect(420, 605, 150, 35)
    pygame.draw.rect(screen, GREEN, btn_rect, border_radius=6)
    btn_text = small_font.render("重新开始", True, WHITE)
    screen.blit(btn_text, (430, 608))
    return btn_rect

def reset_game():
    global board, game_over, winner
    board = [[None, None, None],[None, None, None],[None, None, None]]
    game_over = False
    winner = None

# ---------------------- 主循环 ----------------------
clock = pygame.time.Clock()
running = True
while running:
    btn_rect = draw_board()
    pygame.display.update()

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        if event.type == pygame.MOUSEBUTTONDOWN:
            mx, my = pygame.mouse.get_pos()
            # 点击重启按钮
            if btn_rect.collidepoint(mx, my):
                reset_game()
                continue
            if game_over:
                continue
            # 玩家落子
            row = my // CELL_SIZE
            col = mx // CELL_SIZE
            if 0<=row<3 and 0<=col<3 and board[row][col] is None:
                board[row][col] = player
                if check_win(board, player):
                    winner = player
                    game_over = True
                elif is_board_full(board):
                    game_over = True
                else:
                    # AI走棋
                    ai_move = get_ai_move()
                    if ai_move:
                        ai_r, ai_c = ai_move
                        board[ai_r][ai_c] = ai
                        if check_win(board, ai):
                            winner = ai
                            game_over = True
                        elif is_board_full(board):
                            game_over = True
    clock.tick(60)

pygame.quit()
sys.exit()