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"""
    board = np.array(board)
    
    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."
    
    # Convert Dataframe to numpy array for game logic
    board_array = np.array([[0 if cell == "⚪" else 1 if cell == "🔴" else 2 for cell in row] 
                           for row in board.values.tolist()])
    
    # Player move
    board_array = make_move(board_array, col, 1)
    if check_winner(board_array):
        state["game_over"] = True
        return format_board(board_array), "You win! 🎉"
    
    # AI move
    config = Configuration({"rows": 6, "columns": 7, "inarow": 4})
    obs = Observation({"board": board_array.flatten().tolist(), "mark": 2})
    
    ai_col = my_agent(obs, config)
    board_array = make_move(board_array, ai_col, 2)
    
    if check_winner(board_array):
        state["game_over"] = True
        return format_board(board_array), "AI wins! 🤖"
    
    return format_board(board_array), "Your turn!"

def create_ui():
    with gr.Blocks(css=css) as demo:
        gr.Markdown("# Play Connect Four Against AI")
        gr.Markdown("You are Red (🔴), AI is Yellow (🟡)")
        
        state = gr.State({"game_over": False})
        
        with gr.Row():
            # Add column buttons at the top
            buttons = [gr.Button(str(i)) for i in range(7)]
        
        # Board display below the buttons
        board = gr.Dataframe(
            value=format_board(initialize_game()),
            interactive=False,
            show_label=False,
            headers=None,
            wrap=True,
            elem_id="board"
        )
        
        message = gr.Textbox(value="Your turn!", label="Status")
        new_game = gr.Button("New Game")
        
        def reset_game():
            state = {"game_over": False}
            return format_board(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

# CSS definition
css = """
#board {
    max-width: 400px;
    margin: 0 auto;
}
#board td {
    text-align: center;
    font-size: 24px;
    padding: 8px;
}
"""

demo = create_ui()
demo.launch()