import os
from typing import Optional
import gradio as gr
import requests
from smolagents import CodeAgent, Tool
from smolagents.models import HfApiModel
from smolagents.monitoring import LogLevel
from gradio import ChatMessage

DEFAULT_MODEL = "Qwen/Qwen2.5-Coder-32B-Instruct"
HF_API_TOKEN = os.getenv("HF_TOKEN")

# Tool descriptions for the UI
TOOL_DESCRIPTIONS = {
    "Hub Collections": "Add tool collections from Hugging Face Hub.",
    "Spaces": "Add tools from Hugging Face Spaces.",
}


def search_spaces(query, limit=1):
    """
    Search for Hugging Face Spaces using the API.
    Returns the first result or None if no results.
    """
    try:
        url = f"https://huggingface.co/api/spaces?search={query}&limit={limit}"
        response = requests.get(
            url, headers={"Authorization": f"Bearer {HF_API_TOKEN}"}
        )
        response.raise_for_status()

        spaces = response.json()
        if not spaces:
            return None

        # Get the first space
        space = spaces[0]
        space_id = space["id"]

        # Extract title and description
        title = space_id.split("/")[-1]  # Default to the last part of the ID
        description = f"Tool from {space_id}"

        # Try to get title from different possible locations
        if "title" in space:
            title = space["title"]
        elif "cardData" in space and "title" in space["cardData"]:
            title = space["cardData"]["title"]

        # Try to get description from different possible locations
        if "description" in space:
            description = space["description"]
        elif "cardData" in space and "description" in space["cardData"]:
            description = space["cardData"]["description"]

        return {
            "id": space_id,
            "title": title,
            "description": description,
        }
    except Exception as e:
        print(f"Error searching spaces: {e}")
        return None


def get_space_metadata(space_id):
    """
    Get metadata for a specific Hugging Face Space.
    """
    try:
        url = f"https://huggingface.co/api/spaces/{space_id}"
        response = requests.get(
            url, headers={"Authorization": f"Bearer {HF_API_TOKEN}"}
        )
        response.raise_for_status()

        space = response.json()

        # Extract title and description from the space data
        # The structure can vary, so we need to handle different cases
        title = space_id
        description = f"Tool from {space_id}"

        # Try to get title from different possible locations
        if "title" in space:
            title = space["title"]
        elif "cardData" in space and "title" in space["cardData"]:
            title = space["cardData"]["title"]
        else:
            # Use the last part of the space_id as a fallback title
            title = space_id.split("/")[-1]

        # Try to get description from different possible locations
        if "description" in space:
            description = space["description"]
        elif "cardData" in space and "description" in space["cardData"]:
            description = space["cardData"]["description"]

        return {
            "id": space_id,
            "title": title,
            "description": description,
        }
    except Exception as e:
        print(f"Error getting space metadata: {e}")
        return None


def create_agent(model_name, space_tools=None):
    """
    Create a CodeAgent with the specified model and tools.
    """
    if not space_tools:
        space_tools = []

    try:
        # Convert space tools to Tool objects
        tools = []
        for tool_info in space_tools:
            space_id = tool_info["id"]
            tool = Tool.from_space(
                space_id,
                name=tool_info.get("name", space_id),
                description=tool_info.get("description", ""),
            )
            tools.append(tool)

        # Initialize the HfApiModel with the model name
        model = HfApiModel(model_id=model_name, token=HF_API_TOKEN)

        # Create the agent with the tools and additional imports
        agent = CodeAgent(
            tools=tools,
            model=model,
            additional_authorized_imports=["PIL", "requests"],
            verbosity_level=LogLevel.DEBUG,  # Set higher verbosity for detailed logs
        )
        print(f"Agent created successfully with {len(tools)} tools")
        return agent
    except Exception as e:
        print(f"Error creating agent: {e}")

        # Try with a fallback model if the specified one fails
        try:
            print("Trying fallback model...")
            fallback_model = HfApiModel(
                model_id="Qwen/Qwen2.5-Coder-7B-Instruct", token=HF_API_TOKEN
            )

            agent = CodeAgent(
                tools=tools,
                model=fallback_model,
                additional_authorized_imports=["PIL", "requests"],
                verbosity_level=LogLevel.DEBUG,  # Set higher verbosity for detailed logs
            )
            print("Agent created successfully with fallback model")
            return agent
        except Exception as e:
            print(f"Error creating agent: {e}")
            return None


# Event handler functions
def on_search_spaces(query):
    if not query:
        return "Please enter a search term.", "", "", ""

    try:
        space_info = search_spaces(query)
        if space_info is None:
            return "No spaces found.", "", "", ""

        # Format the results as markdown
        results_md = "### Search Results:\n"
        results_md += f"- ID: `{space_info['id']}`\n"
        results_md += f"- Title: {space_info['title']}\n"
        results_md += f"- Description: {space_info['description']}\n"

        # Return values to update the UI
        return (
            results_md,
            space_info["id"],
            space_info["title"],
            space_info["description"],
        )
    except Exception as e:
        print(f"Error in search: {e}")
        return f"Error: {str(e)}", "", "", ""


def on_validate_space(space_id):
    if not space_id:
        return "Please enter a space ID or search term.", "", ""

    try:
        # First try to get metadata directly if it's a valid space ID
        space_info = get_space_metadata(space_id)

        # If not found, try to search for it
        if space_info is None:
            # Try to search for the space using the ID as a search term
            space_info = search_spaces(space_id)
            if space_info is None:
                return f"No spaces found for '{space_id}'.", "", ""

            # Format search result as markdown
            result_md = f"### Found Space via Search:\n"
            result_md += f"- ID: `{space_info['id']}`\n"
            result_md += f"- Title: {space_info['title']}\n"
            result_md += f"- Description: {space_info['description']}\n"

            return (
                result_md,
                space_info["title"],
                space_info["description"],
            )

        # Format direct match as markdown
        result_md = f"### Space Validated Successfully:\n"
        result_md += f"- ID: `{space_info['id']}`\n"
        result_md += f"- Title: {space_info['title']}\n"
        result_md += f"- Description: {space_info['description']}\n"

        return (
            result_md,
            space_info["title"],
            space_info["description"],
        )
    except Exception as e:
        print(f"Error validating space: {e}")
        return f"Error: {str(e)}", "", ""


def on_add_tool(space_id, space_name, space_description, current_tools):
    if not space_id:
        return (
            current_tools,
            "Please enter a space ID.",
        )

    # Check if this tool is already added
    for tool in current_tools:
        if tool["id"] == space_id:
            return (
                current_tools,
                f"Tool '{space_id}' is already added.",
            )

    # Add the new tool
    new_tool = {
        "id": space_id,
        "name": space_name if space_name else space_id,
        "description": space_description if space_description else "No description",
    }

    updated_tools = current_tools + [new_tool]

    # Format the tools as markdown
    tools_md = "### Added Tools:\n"
    for i, tool in enumerate(updated_tools, 1):
        tools_md += f"{i}. **{tool['name']}** (`{tool['id']}`)\n"
        tools_md += f"   {tool['description']}\n\n"

    return updated_tools, tools_md


def on_create_agent(model, space_tools):
    if not space_tools:
        return (
            None,
            [],
            "",
            "Please add at least one tool before creating an agent.",
            "No agent created yet.",
        )

    try:
        # Create the agent
        agent = create_agent(model, space_tools)

        if agent is None:
            return (
                None,
                [],
                "",
                "Failed to create agent. Please try again with different tools or model.",
                "No agent created yet.",
            )

        # Format the tools for display
        tools_str = ", ".join(
            [f"{tool['name']} ({tool['id']})" for tool in space_tools]
        )

        # Generate agent status
        agent_status = update_agent_status(agent)

        return (
            agent,
            [],
            "",
            f"✅ Agent created successfully with {model}!\nTools: {tools_str}",
            agent_status,
        )
    except Exception as e:
        print(f"Error creating agent: {e}")
        return None, [], "", f"Error creating agent: {str(e)}", "No agent created yet."


def add_user_message(message, chat_history):
    """Add the user message to the chat history."""
    # For Gradio chatbot with type="messages", we need to use ChatMessage objects
    if not message:
        return "", chat_history

    # Add user message to chat history
    chat_history = chat_history + [ChatMessage(role="user", content=message)]
    return message, chat_history


def stream_to_gradio(
    agent,
    task: str,
    reset_agent_memory: bool = False,
    additional_args: Optional[dict] = None,
):
    """Runs an agent with the given task and streams the messages from the agent as gradio ChatMessages."""

    from smolagents.gradio_ui import pull_messages_from_step, handle_agent_output_types
    from smolagents.agent_types import AgentAudio, AgentImage, AgentText

    for step_log in agent.run(
        task, stream=True, reset=reset_agent_memory, additional_args=additional_args
    ):
        for message in pull_messages_from_step(
            step_log,
        ):
            yield message

    final_answer = step_log  # Last log is the run's final_answer
    final_answer = handle_agent_output_types(final_answer)

    if isinstance(final_answer, AgentImage):
        yield gr.ChatMessage(
            role="assistant",
            content={"path": final_answer.to_string(), "mime_type": "image/png"},
        )
    elif isinstance(final_answer, AgentText) and os.path.exists(
        final_answer.to_string()
    ):
        yield gr.ChatMessage(
            role="assistant",
            content=gr.Image(final_answer.to_string()),
        )
    elif isinstance(final_answer, AgentAudio):
        yield gr.ChatMessage(
            role="assistant",
            content={"path": final_answer.to_string(), "mime_type": "audio/wav"},
        )
    else:
        yield gr.ChatMessage(
            role="assistant", content=f"**Final answer:** {str(final_answer)}"
        )


def stream_agent_response(agent, message, chat_history):
    """Stream the agent's response to the chat history."""
    if not message or agent is None:
        return chat_history

    # First yield the current chat history
    yield chat_history

    try:
        # Stream the agent's response
        for msg in stream_to_gradio(agent, message):
            # Add the message to chat history
            chat_history = chat_history + [msg]
            # Yield updated chat history
            yield chat_history
    except Exception as e:
        # Handle errors
        error_msg = f"Error: {str(e)}"
        chat_history = chat_history + [ChatMessage(role="assistant", content=error_msg)]
        yield chat_history


def on_clear(agent=None):
    """Clear the chat and reset the agent."""
    return (
        agent,
        [],
        "",
        "Agent cleared. Create a new one to continue.",
        "",
        gr.update(interactive=False),
    )


def update_agent_status(agent):
    """Update the agent status display with current information."""
    if agent is None:
        return "No agent created yet. Add a Space tool to get started."

    # Get agent information
    tools = agent.tools if hasattr(agent, "tools") else []
    tool_count = len(tools)

    # Create status message
    status = f"Agent ready with {tool_count} tools"
    return status


# Create the Gradio app
with gr.Blocks(title="AI Agent Builder") as app:
    gr.Markdown("# AI Agent Builder with SmolaGents")
    gr.Markdown("Build your own AI agent by selecting tools from Hugging Face Spaces.")

    # Agent state
    agent_state = gr.State(None)
    last_message = gr.State("")
    space_tools_state = gr.State([])

    # Message store for preserving user message
    msg_store = gr.State("")

    with gr.Row():
        # Left sidebar for tool configuration
        with gr.Column(scale=1):
            gr.Markdown("## Tool Configuration")
            gr.Markdown("Add multiple Hugging Face Spaces as tools for your agent:")

            # Hidden model input with default value
            model_input = gr.Textbox(
                value=DEFAULT_MODEL,
                label="Model",
                visible=False,
            )

            # Space tool input
            with gr.Group():
                gr.Markdown("### Add Space as Tool")
                space_tool_input = gr.Textbox(
                    label="Space ID or Search Term",
                    placeholder=("Enter a Space ID or search term"),
                    info="Enter a Space ID (username/space-name) or search term",
                )
                space_name_input = gr.Textbox(
                    label="Tool Name (optional)",
                    placeholder="Enter a name for this tool",
                )
                space_description_input = gr.Textbox(
                    label="Tool Description (optional)",
                    placeholder="Enter a description for this tool",
                    lines=2,
                )
                add_tool_button = gr.Button("Add Tool", variant="primary")

            # Display added tools
            gr.Markdown("### Added Tools")
            tools_display = gr.Markdown(
                "No tools added yet. Add at least one tool before creating an agent."
            )

            # Create agent button
            create_button = gr.Button(
                "Create Agent with Selected Tools", variant="secondary", size="lg"
            )

            # Status message
            status_msg = gr.Markdown("")

            # Agent status display
            agent_status = gr.Markdown("No agent created yet.")

        # Main content area
        with gr.Column(scale=2):
            # Chat interface for the agent
            chatbot = gr.Chatbot(
                label="Agent Chat",
                height=600,
                show_copy_button=True,
                avatar_images=("👤", "🤖"),
                type="messages",  # Use messages type for ChatMessage objects
            )
            msg = gr.Textbox(
                label="Your message",
                placeholder="Type a message to your agent...",
                interactive=True,
            )

            with gr.Row():
                with gr.Column(scale=1, min_width=60):
                    clear = gr.Button("🗑️", scale=1)
                with gr.Column(scale=8):
                    # Empty column for spacing
                    pass

    # Connect event handlers
    # Connect the space_tool_input submit event to the validation handler
    space_tool_input.submit(
        on_validate_space,
        inputs=[space_tool_input],
        outputs=[status_msg, space_name_input, space_description_input],
    )

    # Connect the add tool button
    add_tool_button.click(
        on_add_tool,
        inputs=[
            space_tool_input,
            space_name_input,
            space_description_input,
            space_tools_state,
        ],
        outputs=[space_tools_state, tools_display],
    )

    # Connect the create button to the handler
    create_button.click(
        on_create_agent,
        inputs=[model_input, space_tools_state],
        outputs=[agent_state, chatbot, msg, status_msg, agent_status],
    )

    # Connect the message input to the chain of handlers
    msg.submit(
        lambda message: (message, message, ""),  # Store message and clear input
        inputs=[msg],
        outputs=[msg_store, msg, msg],
        queue=False,
    ).then(
        add_user_message,  # Add user message to chat
        inputs=[msg_store, chatbot],
        outputs=[msg_store, chatbot],
        queue=False,
    ).then(
        stream_agent_response,  # Generate and stream response
        inputs=[agent_state, msg_store, chatbot],
        outputs=chatbot,
        queue=True,
    )


if __name__ == "__main__":
    app.queue().launch()