import gradio as gr
from PIL import Image, ImageDraw, ImageFont
import random
import time

def choose_move(board, player):
    valid = [(r, c) for r in range(8) for c in range(8) if get_flips(board, r, c, player)]
    return random.choice(valid) if valid else None

def initialize_board():
    b = [[0]*8 for _ in range(8)]
    b[3][3], b[4][4] = 1, 1
    b[3][4], b[4][3] = -1, -1
    return b

DIRECTIONS = [(-1,-1),(-1,0),(-1,1),(0,-1),(0,1),(1,-1),(1,0),(1,1)]

def get_flips(board, r, c, p):
    if board[r][c] != 0: return []
    flips = []
    for dr, dc in DIRECTIONS:
        rr, cc = r+dr, c+dc; buf = []
        while 0<=rr<8 and 0<=cc<8 and board[rr][cc] == -p:
            buf.append((rr, cc)); rr+=dr; cc+=dc
        if buf and 0<=rr<8 and 0<=cc<8 and board[rr][cc] == p:
            flips += buf
    return flips

def apply_move(board, r, c, p):
    f = get_flips(board, r, c, p)
    if not f: return False
    board[r][c] = p
    for rr, cc in f: board[rr][cc] = p
    return True

def count_score(board):
    b = sum(cell==-1 for row in board for cell in row)
    w = sum(cell==1  for row in board for cell in row)
    return b, w

def board_to_image(board, last_user, last_ai, player):
    size = 360; cell = size//8
    img = Image.new('RGB', (size, size+cell), 'darkgreen')
    draw = ImageDraw.Draw(img)
    # Scoreboard
    b, w = count_score(board)
    draw.rectangle([0,0,size,cell], fill='navy')
    font = ImageFont.load_default()
    draw.text((10,2), f"Black: {b}", font=font, fill='yellow')
    draw.text((size-80,2), f"White: {w}", font=font, fill='yellow')
    # Board
    for r in range(8):
        for c in range(8):
            x0, y0 = c*cell, cell + r*cell
            x1, y1 = x0+cell, y0+cell
            draw.rectangle([x0,y0,x1,y1], outline='black')
            # possible moves
            if board[r][c]==0 and get_flips(board, r, c, player):
                draw.ellipse([x0+cell*0.4, y0+cell*0.4, x0+cell*0.6, y0+cell*0.6], fill='yellow')
            if board[r][c]==1:
                draw.ellipse([x0+4, y0+4, x1-4, y1-4], fill='white')
            if board[r][c]==-1:
                draw.ellipse([x0+4, y0+4, x1-4, y1-4], fill='black')
    # markers
    if last_user:
        ur, uc = last_user
        x0, y0 = uc*cell, cell + ur*cell; x1, y1 = x0+cell, y0+cell
        draw.rectangle([x0,y0,x1,y1], outline='red', width=4)
    if last_ai:
        ar, ac = last_ai
        x0, y0 = ac*cell, cell + ar*cell; x1, y1 = x0+cell, y0+cell
        draw.rectangle([x0,y0,x1,y1], outline='blue', width=4)
    return img

def click_handler(evt: gr.SelectData, state):
    board, player, last_user, last_ai = state
    x, y = evt.index; cell = 360//8
    c, r = int(x//cell), int((y-cell)//cell)
    # invalid
    if not (0<=r<8 and 0<=c<8):
        yield board_to_image(board, last_user, last_ai, player), state
        return
    # user move
    if apply_move(board, r, c, player):
        last_user = (r, c); player = -player
        yield board_to_image(board, last_user, last_ai, player), state
    else:
        yield board_to_image(board, last_user, last_ai, player), state
        return
    # AI thinking indicator
    yield board_to_image(board, last_user, last_ai, player), (board, player, last_user, last_ai)
    time.sleep(2)
    # AI move
    ai_mv = choose_move(board, player)
    if ai_mv:
        apply_move(board, ai_mv[0], ai_mv[1], player)
        last_ai = ai_mv; player = -player
        yield board_to_image(board, last_user, last_ai, player), (board, player, last_user, last_ai)
    else:
        yield board_to_image(board, last_user, last_ai, player), (board, player, last_user, last_ai)

with gr.Blocks() as demo:
    state = gr.State((initialize_board(), -1, None, None))
    board_img = gr.Image(value=board_to_image(initialize_board(), None, None, -1), interactive=True)
    board_img.select(click_handler, inputs=[state], outputs=[board_img, state])
    demo.launch()