import pygame

pygame.init()
WIDTH = 640
HEIGHT = 640
screen = pygame.display.set_mode((WIDTH, HEIGHT))
pygame.display.set_caption("Gobang")
clock = pygame.time.Clock()

# Colors
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)
BOARD_COLOR = (222, 179, 111)
RED = (220, 30, 30)

# Font
try:
    font = pygame.font.SysFont("Arial", 36)
except:
    font = pygame.font.Font(None, 36)

GRID_SIZE = 40
OFFSET = 40
BOARD_LINE = 15

# Board data 0:empty 1:black 2:white
board = [[0 for _ in range(BOARD_LINE)] for _ in range(BOARD_LINE)]
turn = 1  # 1 black first
game_over = False
winner = 0

def check_win(x, y, player):
    # Check four directions
    directions = [(1,0), (0,1), (1,1), (1,-1)]
    for dx, dy in directions:
        count = 1
        # Forward
        cx, cy = x+dx, y+dy
        while 0 <= cx < BOARD_LINE and 0 <= cy < BOARD_LINE and board[cy][cx] == player:
            count += 1
            cx += dx
            cy += dy
        # Backward
        cx, cy = x-dx, y-dy
        while 0 <= cx < BOARD_LINE and 0 <= cy < BOARD_LINE and board[cy][cx] == player:
            count += 1
            cx -= dx
            cy -= dy
        if count >= 5:
            return True
    return False

running = True
while running:
    screen.fill(BOARD_COLOR)
    mouse_x, mouse_y = pygame.mouse.get_pos()

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False
        if event.type == pygame.MOUSEBUTTONDOWN and not game_over:
            mx, my = event.pos
            # Compute grid position
            gx = round((mx - OFFSET) / GRID_SIZE)
            gy = round((my - OFFSET) / GRID_SIZE)
            if 0 <= gx < BOARD_LINE and 0 <= gy < BOARD_LINE:
                if board[gy][gx] == 0:
                    board[gy][gx] = turn
                    if check_win(gx, gy, turn):
                        game_over = True
                        winner = turn
                    turn = 2 if turn == 1 else 1

    # Draw board lines
    for i in range(BOARD_LINE):
        # Horizontal
        pygame.draw.line(screen, BLACK, (OFFSET, OFFSET + i*GRID_SIZE),
                         (OFFSET + (BOARD_LINE-1)*GRID_SIZE, OFFSET + i*GRID_SIZE), 2)
        # Vertical
        pygame.draw.line(screen, BLACK, (OFFSET + i*GRID_SIZE, OFFSET),
                         (OFFSET + i*GRID_SIZE, OFFSET + (BOARD_LINE-1)*GRID_SIZE), 2)

    # Draw stones
    for y in range(BOARD_LINE):
        for x in range(BOARD_LINE):
            px = OFFSET + x * GRID_SIZE
            py = OFFSET + y * GRID_SIZE
            if board[y][x] == 1:
                pygame.draw.circle(screen, BLACK, (px, py), 16)
            elif board[y][x] == 2:
                pygame.draw.circle(screen, WHITE, (px, py), 16)
                pygame.draw.circle(screen, BLACK, (px, py), 16, 2)

    # Game info
    if not game_over:
        text = font.render("Black Turn" if turn == 1 else "White Turn", True, BLACK)
    else:
        if winner == 1:
            text = font.render("Black Wins!", True, RED)
        else:
            text = font.render("White Wins!", True, RED)
    screen.blit(text, (20, 10))

    pygame.display.flip()
    clock.tick(60)

pygame.quit()