import gradio as gr
import os
import shutil
from gradio_client import Client, handle_file
from smolagents import Tool, CodeAgent, HfApiModel

# import spaces - if using ZeroGPU

# Define tools from Spaces
spaces = [
    {"repo_id": "black-forest-labs/FLUX.1-schnell", 
     "name": "image_generator", 
     "description": "Generate an image from a prompt"},
    
    {"repo_id": "jamesliu1217/EasyControl_Ghibli", 
     "name": "Ghibli_style_Image_control", 
     "description": "Create Ghibli style image"},
]

tools = []
for space in spaces:
    # Access repo_id, name, and description
    repo_id = space['repo_id']
    name = space.get('name', repo_id)  # Use repo_id as name if not specified
    description = space.get('description', '')  # Use empty string if not specified

    # Create Tool instance
    tool = Tool.from_space(repo_id, name=name, description=description)
    tools.append(tool)

# Define a custom tool
class CustomTool(Tool):
    name = "custom_tool"
    description = "A custom tool that processes input text"
    inputs = {"input": {"type": "string", "description": "Some input text to process"}}
    output_type = "string"
    def forward(self, input: str):
        return f"Processed: {input}"

tools.append(CustomTool())


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

# Create the agent
agent = CodeAgent(tools=tools, model=model)

# Function to run the agent and return the image path
def generate_and_transform(prompt):
    result = agent.run(prompt)
    
    if isinstance(result, str):  # Assuming result is a file path
        # Copy the temporary file to a permanent location
        permanent_path = "ghibli_output.webp"
        shutil.copy(result, permanent_path)
        return permanent_path
    else:
        raise ValueError("Unexpected result type from agent")

# Gradio interface function
def gradio_interface(prompt):
    try:
        image_path = generate_and_transform(prompt)
        return image_path
    except Exception as e:
        return str(e)

# Create the Gradio app
with gr.Blocks() as app:
    gr.Markdown("### Smolagent Image Generator with Ghibli Style")
    with gr.Row():
        prompt_input = gr.Textbox(label="Enter your prompt", placeholder="e.g., Generate an image of a dog and then make a Ghibli style version of that image")
        submit_button = gr.Button("Generate")
    output_image = gr.Image(label="Generated Image")
    download_button = gr.File(label="Download Image")

    # Connect the button to the function
    def on_submit(prompt):
        image_path = gradio_interface(prompt)
        return image_path, image_path

    submit_button.click(on_submit, inputs=prompt_input, outputs=[output_image, download_button])

# Launch the app
app.launch()