import os
import warnings

from dotenv import load_dotenv
from huggingface_hub import get_token
from mcp import StdioServerParameters
from smolagents import (
    CodeAgent,
    #    DuckDuckGoSearchTool,
    FinalAnswerTool,
    InferenceClientModel,
    MCPClient,
    Tool,
    #    VisitWebpageTool,
)

from my_ui import CustomGradioUI

# Load environment variables from .env file
load_dotenv()

# Configuration
HF_MODEL_ID = os.getenv("HF_MODEL_ID", "Qwen/Qwen2.5-Coder-32B-Instruct")
# GROQ_API_KEY = os.getenv("GROQ_API_KEY", "")
TAVILY_API_KEY = os.getenv("TAVILY_API_KEY", "")

# Suppress OAuth warning when running locally (OAuth only works on HF Spaces)
warnings.filterwarnings(
    "ignore", message=".*Gradio does not support OAuth features outside of a Space environment.*"
)

# smolagents Toolbox
# web_search = DuckDuckGoSearchTool(max_results=5, rate_limit=2.0)
# visit_webpage = VisitWebpageTool()
final_answer = FinalAnswerTool()

# Import tool from Hub
image_generation_tool = Tool.from_space(
    space_id="black-forest-labs/FLUX.1-schnell",
    name="image_generator",
    description="Generates an image following your prompt. Returns a PIL Image.",
    api_name="/infer",
)

# image_analysis_tool = Tool.from_space(
#     space_id="fffiloni/Sa2VA-simple-demo",
#     name="image_analyzer",
#     description=(
#         "Perform image-based visual question answering and segmentation. "
#         "This function takes an image and a text prompt (instruction) as input, "
#         "processes the image with a multimodal model, and returns a textual answer."
#     ),
#     api_name="/image_vision",
# )


# MCP Servers
def get_upload_files_to_gradio_server():
    """Get the upload files to Gradio MCP server configuration."""
    return StdioServerParameters(
        command="uvx",  # Using uvx ensures dependencies are available
        args=[
            "--from",
            "gradio[mcp]",
            "gradio",
            "upload-mcp",
            "https://fffiloni-sa2va-simple-demo.hf.space/",
            "./data",
        ],
        env={"UV_PYTHON": "3.12", **os.environ},
    )


# groq_compount_beta = StdioServerParameters(
#     command="uvx",  # Using uvx ensures dependencies are available
#     args=[
#         "--from",
#         "git+https://github.com/groq/groq-mcp-server",
#         "groq-mcp"
#     ],
#     env={"GROQ_API_KEY": GROQ_API_KEY,
#          "BASE_OUTPUT_PATH": "./data"  # Optional: Where to save generated files
#     },
# )


def get_image_analysis_mcp():
    """Get the image analysis MCP configuration."""
    return {
        "url": "https://fffiloni-sa2va-simple-demo.hf.space/gradio_api/mcp/",
        "transport": "streamable-http",
    }


def get_tavily_search_mcp(api_key):
    """Get the Tavily search MCP configuration with the provided API key."""
    return {
        "url": f"https://mcp.tavily.com/mcp/?tavilyApiKey={api_key}",
        "transport": "streamable-http",
    }


def create_agent(tavily_api_key=None, oauth_token=None):
    """Create an agent with the specified Tavily API key and OAuth token.

    Args:
        tavily_api_key: Optional Tavily API key for web search
        oauth_token: OAuth token from Gradio (gr.OAuthToken | None)
    """
    # Use provided API key or fall back to environment variable
    api_key = tavily_api_key or TAVILY_API_KEY

    # Build list of MCP servers - only include Tavily if API key is provided
    mcp_servers = [
        get_upload_files_to_gradio_server(),
        get_image_analysis_mcp(),
    ]

    # Only add Tavily search if we have an API key
    if api_key:
        mcp_servers.append(get_tavily_search_mcp(api_key))

    # Create MCP client with the appropriate servers
    mcp_client = MCPClient(mcp_servers, structured_output=False)

    # Create model with user-specified parameters
    # Priority: OAuth token from Gradio > HF_TOKEN env var > get_token() from CLI login
    if oauth_token is not None:
        # On HF Spaces, use the OAuth token provided by Gradio
        hf_token = oauth_token.token if hasattr(oauth_token, 'token') else oauth_token
    else:
        # Fallback to environment variable or CLI login
        hf_token = os.getenv("HF_TOKEN") or get_token()

    model = InferenceClientModel(
        max_tokens=2096,
        temperature=0.5,
        top_p=0.95,
        model_id=HF_MODEL_ID,
        custom_role_conversions=None,
        token=hf_token,
    )

    # Load instructions from file.  This text will be appended to the system prompt.
    with open("instructions.txt") as f:
        instructions = f.read()

    # Initialize the agent with the configured model and tools
    # Use the global mcp_client's tools instead of creating a new context
    agent = CodeAgent(
        model=model,
        additional_authorized_imports=["*"],
        tools=[
            image_generation_tool,
            *mcp_client.get_tools(),  # Use tools from the MCP client
            # web_search,
            # visit_webpage,
            final_answer,
        ],
        instructions=instructions,
        add_base_tools=False,
        max_steps=10,
        verbosity_level=1,
        planning_interval=3,
    )

    # Print the names and descriptions of the added tools
    print("--- Active Tools ---")
    for tool_name, tool in agent.tools.items():
        print(f"Tool Name: {tool_name}")
        print(f"Description: {tool.description}")
        print("-" * 20)

    return agent


# Example prompts for users
EXAMPLES = [
    "Search the web for the latest developments in quantum computing and summarize the key findings",
    "Generate an image of a futuristic city with flying cars and neon lights",
    "What are the current trends in artificial intelligence research?",
]

gradio_ui = CustomGradioUI(
    create_agent,  # Pass the factory function, not the agent instance
    file_upload_folder="./data",
    reset_agent_memory=True,
    allowed_file_types=[
        ".pdf",
        ".docx",
        ".txt",
        ".md",
        ".json",
        ".csv",
        ".png",
        ".jpeg",
        ".jpg",
    ],  # Customize this list!
    examples=EXAMPLES,
)

if __name__ == "__main__":
    gradio_ui.launch(share=False)