import gradio as gr
from smolagents import CodeAgent,HfApiModel, DuckDuckGoSearchTool, tool
from smolagents.agents import ActionStep
from time import sleep
import helium
from PIL import Image
from io import BytesIO
import os

# Define a function to search images and return the image path for display
@tool
def search_images_for_presentation(query: str) -> str:
    """
    Searches for images related to the given query on the web (using DuckDuckGo search).
    Args:
        query: The query to search for images or logos.
    """
    helium.go_to("https://duckduckgo.com/")
    search_box = helium.find(S("input[type='text']"))
    search_box.write(query)
    search_box.press(Keys.ENTER)

    # Wait for search results to load
    sleep(3)

    # Scroll down and search for images related to the query
    search_item_ctrl_f("image", nth_result=1)
    
    # Simulate an image being found (as example)
    image_path = "img/some_image.png"  # Replace with actual image capture logic
    return image_path


@tool
def save_image_from_browser(image_url: str) -> str:
    """
    Downloads and saves the image from the given URL.
    Args:
        image_url: The URL of the image to be saved.
    Returns:
        str: The path to the saved image.
    """
    # Download image from the URL
    try:
        response = requests.get(image_url)
        response.raise_for_status()
        
        # Save the image locally (you can adjust the save path as needed)
        image_path = f"img/{image_url.split('/')[-1]}"
        with open(image_path, "wb") as file:
            file.write(response.content)
        
        return image_path
    except requests.exceptions.RequestException as e:
        print(f"Error downloading image: {e}")
        return None



# Initialize agent (same as before)
def initialize_agent(model):
    return CodeAgent(
        tools=[DuckDuckGoSearchTool(), save_image_from_browser, search_images_for_presentation],
        model=model,
        max_steps=20,
        verbosity_level=2,
    )

find_images_for_presentation_prompt = "Find images related to presentations, such as slides, logos, or visuals. Search for the following: "
find_images_for_presentation_prompt = "Provide image URLs related to the following for a presentation: "

def run_agent(query: str):
    model = HfApiModel(model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud/')
    agent = initialize_agent(model)
    result = agent.run(find_images_for_presentation_prompt + query)
    return result


# Gradio interface to show images
def gradio_interface(query: str):
    # Call the agent to run the query and return the image
    image_paths = run_agent(query)
    images = [Image.open(image_path) for image_path in image_paths]
    
    # Return the image to Gradio for display
    return images


# # Set up Gradio interface
# gr.Interface(
#     fn=gradio_interface,
#     inputs="text",  # Text input for the search query
#     outputs=gr.Gallery(label="Found Images",columns=3),  # Image output to display the fetched image
#     title="Figure, Image & Logo Finder",
#     description="Enter a query to search for relevant images, logos, or figures for your presentation."
# ).launch()