import gradio as gr
from smolagents import CodeAgent, HfApiModel, Tool
from PIL import Image
import os

# Define the image generation tool
image_generation_tool = Tool.from_space(
    "black-forest-labs/FLUX.1-schnell",
    name="image_generator",
    description="Generate an image from a text prompt."
)

# Initialize the model
model = HfApiModel("Qwen/Qwen2.5-Coder-32B-Instruct")  # Replace with your model

# Create the Smol Agent
agent = CodeAgent(tools=[image_generation_tool], model=model)

# Define the chatbot function
def chat_with_agent(user_input, history):
    if history is None:
        history = []
    
    # Run the agent and debug the output
    response = agent.run(user_input)
    print(f"Agent response: {response} (type: {type(response)})")  # Debug print
    
    # Initialize outputs
    image_output = None
    response_text = ""
    
    # Handle different response types
    if isinstance(response, str):
        # Check if it’s a file path
        if any(response.endswith(ext) for ext in [".webp", ".png", ".jpg", ".jpeg"]) and os.path.exists(response):
            try:
                image_output = Image.open(response)
                response_text = "Image generated! See below."
            except Exception as e:
                response_text = f"Error loading image: {str(e)}"
        else:
            # Treat as plain text
            response_text = response
    elif isinstance(response, int):  # Handle unexpected integer response
        response_text = f"Agent returned a number ({response}), which I can’t display directly. Try asking differently!"
    else:
        # Fallback for unexpected types
        response_text = f"Unexpected response: {str(response)}. I’ll try to handle it better next time!"
    
    # Append to history as a valid [user, assistant] pair
    history.append([user_input, response_text])
    return history, "", image_output

# Create the Gradio interface
with gr.Blocks(title="Smol Agent Chatbot") as demo:
    gr.Markdown("# Chat with a Smol Agent")
    gr.Markdown("Ask anything—text or image generation supported!")
    
    # Chatbot for text conversation
    chatbot = gr.Chatbot()
    
    # Image output for displaying generated images
    image_display = gr.Image(label="Generated Image")
    
    # Textbox for user input
    user_input = gr.Textbox(placeholder="Type your message here...", label="Your Message")
    
    # Submit button
    submit_btn = gr.Button("Send")
    
    # Connect components
    submit_btn.click(
        fn=chat_with_agent,
        inputs=[user_input, chatbot],
        outputs=[chatbot, user_input, image_display]
    )

# Launch the app
demo.launch(debug=True)  # Enable debug mode to see logs