import gradio as gr
import numpy as np
from game_logic import my_agent, Configuration, Observation

def initialize_game():
    return np.zeros((6, 7), dtype=int)

def make_move(board, column, player):
    """Make a move on the board"""
    for row in range(5, -1, -1):
        if board[row][column] == 0:
            board[row][column] = player
            return board
    return board

def format_board(board):
    """Convert board to player-friendly visualization"""
    symbols = {0: "⚪", 1: "🔴", 2: "🟡"}
    return [[symbols[cell] for cell in row] for row in board]

def check_winner(board):
    """Check if there's a winner"""
    # Convert to numpy array for easier checking
    board = np.array(board)
    
    # Check horizontal, vertical and diagonal wins
    def check_line(line):
        return (len(line) >= 4 and 
                (any(np.all(line[i:i+4] == 1) for i in range(len(line)-3)) or
                 any(np.all(line[i:i+4] == 2) for i in range(len(line)-3))))

    # Horizontal
    for row in board:
        if check_line(row):
            return True

    # Vertical
    for col in board.T:
        if check_line(col):
            return True

    # Diagonals
    for offset in range(-2, 4):
        diag = np.diagonal(board, offset)
        if check_line(diag):
            return True
        diag = np.diagonal(np.fliplr(board), offset)
        if check_line(diag):
            return True

    return False

def play_game(board, col, state):
    if state["game_over"]:
        return board, "Game is over! Click 'New Game' to play again."
    
    # Player move
    board = make_move(board, col, 1)
    if check_winner(board):
        state["game_over"] = True
        return board, "You win! 🎉"
    
    # AI move
    config = Configuration({"rows": 6, "columns": 7, "inarow": 4})
    obs = Observation({"board": board.flatten().tolist(), "mark": 2})
    
    ai_col = my_agent(obs, config)
    board = make_move(board, ai_col, 2)
    
    if check_winner(board):
        state["game_over"] = True
        return board, "AI wins! 🤖"
    
    return board, "Your turn!"

def create_ui():
    with gr.Blocks() as demo:
        gr.Markdown("# Play Connect Four Against AI")
        gr.Markdown("You are Red (🔴), AI is Yellow (🟡)")
        
        state = gr.State({"game_over": False})
        board = gr.DataFrame(
            format_board(initialize_game()),
            headers=False,
            interactive=False,
            show_label=False
        )
        message = gr.Textbox(value="Your turn!", label="Status")
        
        with gr.Row():
            buttons = [gr.Button(str(i)) for i in range(7)]
        
        new_game = gr.Button("New Game")
        
        def reset_game():
            state = {"game_over": False}
            return initialize_game(), "Your turn!", state
        
        new_game.click(
            reset_game,
            outputs=[board, message, state]
        )
        
        for i, button in enumerate(buttons):
            button.click(
                play_game,
                inputs=[board, gr.Number(value=i, visible=False), state],
                outputs=[board, message]
            )
    
    return demo

demo = create_ui()
demo.launch()