import pygame
import sys
import math

# 初始化
pygame.init()

# 常量
WIDTH, HEIGHT = 600, 600
LINE_WIDTH = 10
BOARD_ROWS, BOARD_COLS = 3, 3
SQUARE_SIZE = WIDTH // BOARD_COLS
CIRCLE_RADIUS = SQUARE_SIZE // 3
CIRCLE_WIDTH = 15
CROSS_WIDTH = 20
SPACE = SQUARE_SIZE // 4

# 颜色
BG_COLOR = (28, 170, 156)
LINE_COLOR = (23, 145, 135)
CIRCLE_COLOR = (239, 231, 200)
CROSS_COLOR = (66, 66, 66)

# 屏幕
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Tic-Tac-Toe vs AI")
screen.fill(BG_COLOR)

# 棋盘（0=空，1=玩家，2=AI）
board = [[0 for _ in range(3)] for _ in range(3)]

# 玩家先手
player_turn = True

# 绘制网格
def draw_lines():
    for i in range(1, 3):
        pygame.draw.line(
            screen, LINE_COLOR,
            (0, i * SQUARE_SIZE),
            (WIDTH, i * SQUARE_SIZE),
            LINE_WIDTH
        )
        pygame.draw.line(
            screen, LINE_COLOR,
            (i * SQUARE_SIZE, 0),
            (i * SQUARE_SIZE, HEIGHT),
            LINE_WIDTH
        )

# 绘制棋子
def draw_figures():
    for row in range(3):
        for col in range(3):
            if board[row][col] == 1:
                pygame.draw.line(
                    screen, CROSS_COLOR,
                    (col * SQUARE_SIZE + SPACE, row * SQUARE_SIZE + SQUARE_SIZE - SPACE),
                    (col * SQUARE_SIZE + SQUARE_SIZE - SPACE, row * SQUARE_SIZE + SPACE),
                    CROSS_WIDTH
                )
                pygame.draw.line(
                    screen, CROSS_COLOR,
                    (col * SQUARE_SIZE + SPACE, row * SQUARE_SIZE + SPACE),
                    (col * SQUARE_SIZE + SQUARE_SIZE - SPACE, row * SQUARE_SIZE + SQUARE_SIZE - SPACE),
                    CROSS_WIDTH
                )
            elif board[row][col] == 2:
                pygame.draw.circle(
                    screen, CIRCLE_COLOR,
                    (col * SQUARE_SIZE + SQUARE_SIZE // 2,
                     row * SQUARE_SIZE + SQUARE_SIZE // 2),
                    CIRCLE_RADIUS,
                    CIRCLE_WIDTH
                )

# 检查胜利
def check_winner(player):
    for row in range(3):
        if all(board[row][col] == player for col in range(3)):
            return True
    for col in range(3):
        if all(board[row][col] == player for row in range(3)):
            return True
    if all(board[i][i] == player for i in range(3)):
        return True
    if all(board[i][2 - i] == player for i in range(3)):
        return True
    return False

# 是否平局
def is_board_full():
    return all(board[row][col] != 0 for row in range(3) for col in range(3))

# AI：Minimax
def minimax(depth, is_maximizing):
    if check_winner(2):
        return 1
    if check_winner(1):
        return -1
    if is_board_full():
        return 0

    if is_maximizing:
        best_score = -math.inf
        for row in range(3):
            for col in range(3):
                if board[row][col] == 0:
                    board[row][col] = 2
                    score = minimax(depth + 1, False)
                    board[row][col] = 0
                    best_score = max(score, best_score)
        return best_score
    else:
        best_score = math.inf
        for row in range(3):
            for col in range(3):
                if board[row][col] == 0:
                    board[row][col] = 1
                    score = minimax(depth + 1, True)
                    board[row][col] = 0
                    best_score = min(score, best_score)
        return best_score

# AI 落子
def ai_move():
    best_score = -math.inf
    best_move = None

    for row in range(3):
        for col in range(3):
            if board[row][col] == 0:
                board[row][col] = 2
                score = minimax(0, False)
                board[row][col] = 0
                if score > best_score:
                    best_score = score
                    best_move = (row, col)

    if best_move:
        board[best_move[0]][best_move[1]] = 2

# 重置
def restart():
    global board, player_turn
    board = [[0 for _ in range(3)] for _ in range(3)]
    player_turn = True
    screen.fill(BG_COLOR)
    draw_lines()

draw_lines()

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

        if event.type == pygame.MOUSEBUTTONDOWN and player_turn:
            x, y = event.pos
            row = y // SQUARE_SIZE
            col = x // SQUARE_SIZE

            if board[row][col] == 0:
                board[row][col] = 1
                if not check_winner(1) and not is_board_full():
                    player_turn = False
                    ai_move()
                    player_turn = True

        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_r:
                restart()

    draw_figures()

    if check_winner(1):
        print("🎉 你赢了！")
        player_turn = False
    elif check_winner(2):
        print("🤖 AI 赢了！")
        player_turn = False
    elif is_board_full():
        print("🤝 平局！")

    pygame.display.update()