import os
from dotenv import load_dotenv
import gradio as gr
from langchain_huggingface import HuggingFaceEndpoint

# Load environment variables
load_dotenv()
HF_TOKEN = os.getenv("HF_TOKEN")

# Initialize the HuggingFace inference endpoint
llm = HuggingFaceEndpoint(
    repo_id="mistralai/Mistral-7B-Instruct-v0.3",
    huggingfacehub_api_token=HF_TOKEN.strip(),
    temperature=0.7,
)

# Initialize the game board and state
def initialize_game():
    board = [["" for _ in range(3)] for _ in range(3)]
    current_player = "X"
    status = "Player 1's turn (X)"
    buttons = [gr.Button(value="", elem_classes=["cell-btn"], interactive=True) for _ in range(9)]
    return board, current_player, status, *buttons

# Check for a winner
def check_winner(board):
    for i in range(3):
        if board[i][0] == board[i][1] == board[i][2] and board[i][0] != "":
            return board[i][0]
        if board[0][i] == board[1][i] == board[2][i] and board[0][i] != "":
            return board[0][i]
    if board[0][0] == board[1][1] == board[2][2] and board[0][0] != "":
        return board[0][0]
    if board[0][2] == board[1][1] == board[2][0] and board[0][2] != "":
        return board[0][2]
    return None

# Check for a draw
def check_draw(board):
    return all(cell != "" for row in board for cell in row)

# Minimax algorithm for AI's move
def minimax(board, depth, is_maximizing):
    winner = check_winner(board)
    if winner == "X":
        return -10 + depth
    elif winner == "O":
        return 10 - depth
    elif check_draw(board):
        return 0

    if is_maximizing:
        best = -float('inf')
        for i in range(3):
            for j in range(3):
                if board[i][j] == "":
                    board[i][j] = "O"
                    best = max(best, minimax(board, depth + 1, False))
                    board[i][j] = ""
        return best
    else:
        best = float('inf')
        for i in range(3):
            for j in range(3):
                if board[i][j] == "":
                    board[i][j] = "X"
                    best = min(best, minimax(board, depth + 1, True))
                    board[i][j] = ""
        return best

# Find the best move for AI
def get_best_move(board):
    best_val = -float('inf')
    best_move = (-1, -1)
    for i in range(3):
        for j in range(3):
            if board[i][j] == "":
                board[i][j] = "O"
                move_val = minimax(board, 0, False)
                board[i][j] = ""
                if move_val > best_val:
                    best_move = (i, j)
                    best_val = move_val
    return best_move

# Handle a move
def handle_move(board, current_player, button_idx, game_status):
    if "wins" in game_status or "draw" in game_status:
        buttons = [gr.Button(value=board[i//3][i%3], elem_classes=["cell-btn"], interactive=False) for i in range(9)]
        return board, current_player, game_status, *buttons

    row, col = divmod(button_idx, 3)
    if board[row][col] != "":
        status = f"Invalid move! Player {1 if current_player == 'X' else 2}'s turn ({current_player})"
        buttons = [gr.Button(value=board[i//3][i%3], elem_classes=["cell-btn"]) for i in range(9)]
        return board, current_player, status, *buttons

    board[row][col] = current_player
    winner = check_winner(board)
    if winner:
        status = f"Player {1 if winner == 'X' else 2} ({winner}) wins! 🎉"
        buttons = [gr.Button(value=board[i//3][i%3], elem_classes=["cell-btn"], interactive=False) for i in range(9)]
        return board, current_player, status, *buttons

    if check_draw(board):
        status = "It's a draw! 🤝"
        buttons = [gr.Button(value=board[i//3][i%3], elem_classes=["cell-btn"], interactive=False) for i in range(9)]
        return board, current_player, status, *buttons

    # AI's turn
    if current_player == "X":
        current_player = "O"
        ai_row, ai_col = get_best_move(board)
        board[ai_row][ai_col] = "O"
        winner = check_winner(board)
        if winner:
            status = f"AI ({winner}) wins! 🎉"
            buttons = [gr.Button(value=board[i//3][i%3], elem_classes=["cell-btn"], interactive=False) for i in range(9)]
            return board, current_player, status, *buttons

        if check_draw(board):
            status = "It's a draw! 🤝"
            buttons = [gr.Button(value=board[i//3][i%3], elem_classes=["cell-btn"], interactive=False) for i in range(9)]
            return board, current_player, status, *buttons

        current_player = "X"
        status = f"Player 1's turn (X)"

    buttons = [gr.Button(value=board[i//3][i%3], elem_classes=["cell-btn"]) for i in range(9)]
    return board, current_player, status, *buttons

# Generate a hint using LLM
def get_hint_from_llm(board):
    prompt = f"The current Tic-Tac-Toe board state is {board}. Suggest the best move for Player X with reasoning."
    hint = llm(prompt)
    return hint

# Build the Gradio UI
with gr.Blocks(css=".cell-btn {height: 100px; width: 100px; font-size: 2em; text-align: center;}") as tic_tac_toe:
    gr.Markdown("## Tic-Tac-Toe with AI 🎮")

    # Initialize states
    board_state = gr.State([["" for _ in range(3)] for _ in range(3)])
    current_player = gr.State("X")
    game_status = gr.Textbox(value="Player 1's turn (X)", label="Game Status", interactive=False)

    # Create grid buttons
    buttons = []
    for i in range(3):
        with gr.Row():
            for j in range(3):
                btn = gr.Button(value="", elem_classes=["cell-btn"])
                buttons.append(btn)

    # Hint button
    hint_button = gr.Button("Get Hint")
    hint_display = gr.Textbox(value="", label="Hint", interactive=False)
    hint_button.click(get_hint_from_llm, inputs=[board_state], outputs=[hint_display])

    # Update buttons dynamically on click
    for idx, btn in enumerate(buttons):
        btn.click(
            handle_move,
            inputs=[board_state, current_player, gr.Number(idx, visible=False), game_status],
            outputs=[board_state, current_player, game_status, *buttons],
        )

    # Reset game button
    reset_button = gr.Button("Reset Game")
    reset_button.click(
        initialize_game,
        inputs=[],
        outputs=[board_state, current_player, game_status, *buttons],
    )

tic_tac_toe.launch()