import gradio as gr
import os
from smolagents import CodeAgent, ToolCollection, Tool
from smolagents.models import HfApiModel, LiteLLMModel

# Default model to use
DEFAULT_MODEL = "Qwen/Qwen2.5-Coder-7B-Instruct"

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


# Function to create an agent with selected tools
def create_agent(model_name, hub_tool=None, space_tool=None):
    tools = []

    # Add tool from Hub if provided
    if hub_tool:
        try:
            hub_collection = ToolCollection.from_hub(collection_slug=hub_tool)
            tools.extend(hub_collection.tools)
        except Exception as e:
            print(f"Error loading Hub tool: {e}")

    # Add tool from Space if provided
    if space_tool:
        try:
            space_tool_obj = Tool.from_space(
                space_id=space_tool,
                name=f"space_{space_tool.replace('/', '_')}",
                description=f"Tool from Hugging Face Space: {space_tool}",
            )
            tools.append(space_tool_obj)
        except Exception as e:
            print(f"Error loading Space tool: {e}")

    # Create and return the agent
    try:
        # Try to use HfApiModel first
        model = HfApiModel(model_id=model_name)
        return CodeAgent(tools=tools, model=model)
    except Exception:
        # Fall back to LiteLLMModel if HfApiModel fails
        try:
            model = LiteLLMModel(model_id=model_name)
            return CodeAgent(tools=tools, model=model)
        except Exception as e:
            print(f"Error creating agent: {e}")
            return None


# Main application
def main():
    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 Hub and Spaces."
        )

        with gr.Tabs():
            with gr.TabItem("Build Agent"):
                with gr.Row():
                    with gr.Column(scale=1):
                        # Model selection
                        model_input = gr.Textbox(
                            label="Model Name",
                            placeholder="Enter model name or ID",
                            value=DEFAULT_MODEL,
                        )

                        # Hub tool input
                        hub_tool_input = gr.Textbox(
                            label="Add Tool Collection from Hub (collection slug)",
                            placeholder="e.g., huggingface-tools/diffusion-tools-...",
                        )

                        # Space tool input
                        space_tool_input = gr.Textbox(
                            label="Add Tool from Space (space ID)",
                            placeholder="e.g., black-forest-labs/FLUX.1-schnell",
                        )

                        # Create agent button
                        create_button = gr.Button("Create Agent")

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

                    with gr.Column(scale=2):
                        # Chat interface for the agent
                        chatbot = gr.Chatbot(label="Agent Chat")
                        msg = gr.Textbox(label="Your message")

                        with gr.Row():
                            clear = gr.Button("Clear Chat")
                            regenerate = gr.Button("Regenerate Response")

            with gr.TabItem("Tool Descriptions"):
                tool_descriptions_md = """
## Hugging Face Hub Tool Collections

You can add tool collections from Hugging Face Hub by providing the collection slug.
Example: `huggingface-tools/diffusion-tools-6630bb19a942c2306a2cdb6f`

## Hugging Face Spaces as Tools

You can add tools from Hugging Face Spaces by providing the space ID.
Example: `black-forest-labs/FLUX.1-schnell`

This allows you to use any Gradio app on Hugging Face Spaces as a tool for your agent.
                """

                gr.Markdown(tool_descriptions_md)

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

        # Event handlers
        def on_create_agent(model, hub_tool, space_tool):
            if not model:
                return None, [], "", "⚠️ Please enter a model name."

            if not hub_tool and not space_tool:
                return None, [], "", "⚠️ Please add at least one tool from Hub or Space."

            agent = create_agent(model, hub_tool, space_tool)

            if agent is None:
                return (
                    None,
                    [],
                    "",
                    "❌ Failed to create agent. Check console for details.",
                )

            tools_info = []
            if hub_tool:
                tools_info.append(f"Hub collection: {hub_tool}")
            if space_tool:
                tools_info.append(f"Space: {space_tool}")

            tools_str = " | ".join(tools_info)

            return (
                agent,
                [],
                "",
                f"✅ Agent created successfully with {model}! ({tools_str})",
            )

        create_button.click(
            on_create_agent,
            inputs=[model_input, hub_tool_input, space_tool_input],
            outputs=[agent_state, chatbot, msg, status_msg],
        )

        def on_message(message, chat_history, agent, last_msg):
            if not message:
                return "", chat_history, last_msg

            if agent is None:
                chat_history.append((message, "Please create an agent first."))
                return "", chat_history, last_msg

            try:
                response = agent.run(message, reset=False)
                chat_history.append((message, response))
                return "", chat_history, message
            except Exception as e:
                error_msg = f"Error: {str(e)}"
                chat_history.append((message, error_msg))
                return "", chat_history, message

        msg.submit(
            on_message,
            inputs=[msg, chatbot, agent_state, last_message],
            outputs=[msg, chatbot, last_message],
        )

        def on_regenerate(chat_history, agent, last_msg):
            if not chat_history or not last_msg or agent is None:
                return chat_history, last_msg

            try:
                # Remove the last exchange
                if chat_history:
                    chat_history.pop()

                # Regenerate the response
                response = agent.run(last_msg, reset=False)
                chat_history.append((last_msg, response))
                return chat_history, last_msg
            except Exception as e:
                error_msg = f"Error regenerating response: {str(e)}"
                chat_history.append((last_msg, error_msg))
                return chat_history, last_msg

        regenerate.click(
            on_regenerate,
            inputs=[chatbot, agent_state, last_message],
            outputs=[chatbot, last_message],
        )

        def on_clear():
            return None, [], "", "Agent cleared. Create a new one to continue."

        clear.click(on_clear, outputs=[agent_state, chatbot, last_message, status_msg])

    return app


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