from transformers import load_tool, ReactCodeAgent, HfApiEngine
from PIL import Image
import torch
import numpy as np
import tempfile
import os
import uuid
import gradio as gr


# function to plot and save an AgentImage
def plot_and_save_agent_image(agent_image, save_path=None):
    # Convert AgentImage to a raw PIL Image
    pil_image = agent_image.to_raw()

    # Plot the image using PIL's show method
    pil_image.show()

    # If save_path is provided, save the image
    if save_path:
        pil_image.save(save_path)
        print(f"Image saved to {save_path}")
    else:
        print("No save path provided. Image not saved.")


def generate_prompts_for_object(object_name):
    prompts = {
        "past": f"Show an old version of a {object_name} from its early days.",
        "present": f"Show a modern {object_name} with its current design and technology.",
        "future": f"Show a futuristic version of a {object_name} with advanced features and futuristic design."
    }
    return prompts


# Function to generate the car industry history
def generate_object_history(object_name):
    images = []
    
    # Get prompts for the object
    prompts = generate_prompts_for_object(object_name)
    
    # Generate sequential images and display them
    for time_period, frame in prompts.items():
        print(f"Generating {time_period} frame: {frame}")
        result = agent.run(frame)  # The tool generates the image
        
        # Append the image to the list for GIF creation
        images.append(result.to_raw())  # Ensure we're using raw image for GIF

        # Save each image with the appropriate name (past, present, future)
        image_filename = f"{object_name}_{time_period}.png"
        plot_and_save_agent_image(result, save_path=image_filename)
        
        
    # Create GIF from images
    gif_path = f"{object_name}_evolution.gif"
    images[0].save(
        gif_path, 
        save_all=True, 
        append_images=images[1:], 
        duration=1000,  # Duration in milliseconds for each frame
        loop=0          # Infinite loop
    )
    
    # Return images and GIF path
    return images, gif_path

# Import text-to-image tool from Hub 
# m-ric/text-to-image model generates images based on textual descriptions.
image_generation_tool = load_tool("m-ric/text-to-image", cache=False) #cache=False ensures it fetches the latest tool updates directly from the Hub.

# Import search tool from LangChain
#This tool allows the agent to search for and retrieve information from the web.
from transformers.agents.search import DuckDuckGoSearchTool

search_tool = DuckDuckGoSearchTool()

# Qwen2.5-72B-Instruct is a specific, a LLM fine-tuned for instruction-following tasks.
llm_engine = HfApiEngine("Qwen/Qwen2.5-72B-Instruct") 
# Initialize the agent with both tools
agent = ReactCodeAgent(tools=[image_generation_tool, search_tool], llm_engine=llm_engine)


# Your existing generate_object_history function goes here

# Gradio interface
def create_gradio_interface():
    with gr.Blocks() as demo:
        gr.Markdown("# Object Evolution Generator")
        
        # Add a section for instructions
        gr.Markdown("""
        ## Welcome to the Object Evolution Generator!

        This app allows you to generate visualizations of how an object, like a bicycle or a car, may have evolved over time. 
        It generates images of the object in the past, present, and future based on your input.

        ### How to use:
        - Enter the name of an object (e.g., "bicycle", "car", "phone").
        - Click "Generate Evolution" to generate the evolution of the object across three time periods: past, present, and future.
        - View the generated images and a GIF showing the evolution of the object.

        ### Example:
        Try entering an object name like "car" and see how it has evolved!
        """)

        with gr.Row():
            with gr.Column():
                # Textbox for user to input an object name
                object_name_input = gr.Textbox(label="Enter an object name (e.g., car)", 
                                              placeholder="Enter an object name", 
                                              lines=1)

                # Button to trigger the generation of images and GIF
                generate_button = gr.Button("Generate Evolution")
    
                # Gradio Gallery component to display the images
                image_gallery = gr.Gallery(label="Generated Images", show_label=True, columns=3, rows=1)

                # Output for the generated GIF
                gif_output = gr.Image(label="Generated GIF", show_label=True)
         
        # Set the action when the button is clicked
        generate_button.click(fn=generate_object_history, inputs=[object_name_input], outputs=[image_gallery, gif_output])
    
    return demo

# Launch the Gradio app (permanently)
demo = create_gradio_interface()

# To make it permanent and hosted, we can use Gradio's 'share' argument or host it on a server.
demo.launch(share=True)