import atexit
import os
import threading
from typing import Any

import gradio as gr
import spaces
from huggingface_hub import get_token
from smolagents import CodeAgent, InferenceClientModel, MCPClient


# Your existing MCP sentiment server.
MCP_SERVER_URL = os.getenv(
    "MCP_SERVER_URL",
    "https://zlysunshine-mcp-sentiment.hf.space/gradio_api/mcp/",
)

# This can be changed later through a Space variable.
MODEL_ID = os.getenv(
    "MODEL_ID",
    "Qwen/Qwen3-Next-80B-A3B-Thinking",
)


_agent: CodeAgent | None = None
_mcp_client: MCPClient | None = None
_initialization_lock = threading.Lock()


def create_agent() -> CodeAgent:
    """
    Lazily connect to the remote MCP server and create the agent.

    Returns:
        A CodeAgent configured with the remote MCP tools.
    """
    global _agent
    global _mcp_client

    if _agent is not None:
        return _agent

    with _initialization_lock:
        if _agent is not None:
            return _agent

        # On Hugging Face Spaces, this comes from the HF_TOKEN secret.
        # In Codespaces, it can fall back to the locally saved HF login.
        token = os.getenv("HF_TOKEN") or get_token()

        if not token:
            raise RuntimeError(
                "No Hugging Face token was found. "
                "Set HF_TOKEN or run `hf auth login`."
            )

        client: MCPClient | None = None

        try:
            client = MCPClient(
                {
                    "url": MCP_SERVER_URL,
                    "transport": "streamable-http",
                },
                structured_output=True,
            )

            tools = client.get_tools()

            if not tools:
                raise RuntimeError(
                    "The MCP server connected successfully but returned no tools."
                )

            print("MCP tools discovered:")
            for tool in tools:
                print(f"- {tool.name}: {tool.description}")

            model = InferenceClientModel(
                model_id=MODEL_ID,
                token=token,
                timeout=120,
                max_tokens=1200,
            )

            agent = CodeAgent(
                tools=[*tools],
                model=model,
                max_steps=4,
                additional_authorized_imports=[
                    "json",
                    "ast",
                ],
            )

            _mcp_client = client
            _agent = agent

            return agent

        except Exception:
            if client is not None:
                client.disconnect()
            raise


def close_mcp_connection() -> None:
    """Close the long-lived MCP connection when the app stops."""
    global _mcp_client

    if _mcp_client is not None:
        try:
            _mcp_client.disconnect()
        except Exception as error:
            print(f"Error while closing MCP client: {error}")
        finally:
            _mcp_client = None


atexit.register(close_mcp_connection)


@spaces.GPU(duration=20)
def respond(message: str, history: list[dict[str, Any]]) -> str:
    """
    Answer a user question using tools from the remote MCP server.

    Args:
        message: The user's latest chat message.
        history: Previous Gradio chat messages.

    Returns:
        The agent's final response.
    """
    del history  # The course example treats each request independently.

    cleaned_message = message.strip()

    if not cleaned_message:
        return "Please enter a question."

    try:
        agent = create_agent()
        result = agent.run(cleaned_message)
        return str(result)

    except Exception as error:
        return (
            "The MCP client could not complete the request.\n\n"
            f"Error type: {type(error).__name__}\n"
            f"Details: {error}\n\n"
            "Check that the sentiment MCP server is running and that "
            "the client Space has an HF_TOKEN secret with inference access."
        )


demo = gr.ChatInterface(
    fn=respond,
    title="Gradio MCP Client Agent",
    description=(
        "An AI agent that connects to the "
        "zlysunshine/mcp-sentiment MCP server and uses its tools."
    ),
    examples=[
        (
            "You must use the sentiment_analysis tool to analyze: "
            "'This MCP course is excellent, but deployment was frustrating.'"
        ),
        (
            "Call the sentiment tool for: "
            "'I am happy that the application finally works.'"
        ),
    ],
)


if __name__ == "__main__":
    demo.queue().launch(
        server_name="0.0.0.0",
    )