import asyncio
import os

import chess
import gradio as gr
from gradio_chessboard import Chessboard
from langchain_core.messages import HumanMessage
from utils.helpers import call_agent, create_agent
from utils.tools import create_base_tools, create_mcp_tools

MODEL_DEFAULTS = {
    "anthropic": {
        "model_name": "claude-sonnet-4-20250514",
        "provider": "anthropic",
        "api_key": os.getenv("ANTHROPIC_API_KEY", ""),
    },
    "mistral": {
        "model_name": "mistral-large-latest",
        "provider": "mistralai",
        "api_key": os.getenv("MISTRAL_API_KEY", ""),
    },
    "openai": {
        "model_name": "gpt-4o",
        "provider": "openai",
        "api_key": os.getenv("OPENAI_API_KEY", ""),
    },
}


async def main():
    board = chess.Board()

    base_tools = create_base_tools(board)
    mcp_tools = await create_mcp_tools(
        url="https://agents-mcp-hackathon-chess-mcp-server.hf.space/gradio_api/mcp/sse",
        transport="sse",
    )

    async def chat_entrypoint(
        fen, prompt, messages, model_name, model_provider, api_key
    ):
        """Entrypoint for the chat interaction."""
        board.set_fen(fen)
        messages.append(gr.ChatMessage(role="user", content=prompt))
        yield messages, board.fen()

        real_prompt = HumanMessage(
            content=f"{prompt}\nCurrent board state: {board.fen()}"
        )
        agent = create_agent(
            model_name, model_provider, api_key, base_tools + mcp_tools
        )
        async for messages in call_agent(agent, messages, real_prompt):
            yield messages, board.fen()

    async def move_entrypoint(fen, messages, model_name, model_provider, api_key):
        """Entrypoint for the chess move interaction."""
        board.set_fen(fen)
        messages.append(
            gr.ChatMessage(
                role="user",
                content="Your turn! Make a move.",
            )
        )
        yield messages, board.fen()

        real_prompt = HumanMessage(
            content=f"Make a move with {'white' if board.turn == chess.WHITE else 'black'}, current board state: {fen}. Use can use tools to analyze the position."
        )
        agent = create_agent(
            model_name, model_provider, api_key, base_tools + mcp_tools
        )
        async for messages in call_agent(agent, messages, real_prompt):
            yield messages, board.fen()

    with gr.Blocks(fill_height=True) as chessagent:
        gr.Markdown("# Play Chess with an AI Agent ♔ and see its thoughts 💭")

        with gr.Row():
            with gr.Column(min_width=500):
                with gr.Row():
                    model_provider = gr.Dropdown(
                        choices=["Anthropic", "Mistral", "OpenAI"],
                        value="Anthropic",
                        label="Model Provider",
                        type="value",
                        interactive=True,
                        allow_custom_value=False,
                    )
                    model_name = gr.Textbox(
                        value=MODEL_DEFAULTS["anthropic"]["model_name"],
                        label="Model Name",
                        placeholder="Enter model name",
                        interactive=True,
                    )
                    api_key = gr.Textbox(
                        value=MODEL_DEFAULTS["anthropic"]["api_key"],
                        label="API Key",
                        placeholder="Enter your API key",
                        type="password",
                        interactive=True,
                    )
                    model_provider.change(
                        fn=lambda provider: MODEL_DEFAULTS[provider.lower()][
                            "model_name"
                        ],
                        inputs=model_provider,
                        outputs=model_name,
                    )
                    model_provider.change(
                        fn=lambda provider: MODEL_DEFAULTS[provider.lower()]["api_key"],
                        inputs=model_provider,
                        outputs=api_key,
                    )

                board_component = Chessboard(game_mode=True, label="Chess Board")

                chatbot = gr.Chatbot(
                    type="messages",
                    label="Chess Agent",
                    avatar_images=(
                        "https://chessboardjs.com/img/chesspieces/wikipedia/wK.png",
                        "https://chessboardjs.com/img/chesspieces/wikipedia/bK.png",
                    ),
                    min_height=650,
                    render=False,
                )

                board_component.move(
                    fn=move_entrypoint,
                    inputs=[
                        board_component,
                        chatbot,
                        model_name,
                        model_provider,
                        api_key,
                    ],
                    outputs=[chatbot, board_component],
                )

            with gr.Column():
                chatbot.render()
                input_box = gr.Textbox(lines=1, label="Chat Message")
                input_box.submit(
                    fn=chat_entrypoint,
                    inputs=[
                        board_component,
                        input_box,
                        chatbot,
                        model_name,
                        model_provider,
                        api_key,
                    ],
                    outputs=[chatbot, board_component],
                )
                input_box.submit(lambda: "", None, [input_box], queue=False)

        chessagent.launch()


if __name__ == "__main__":
    asyncio.run(main())