import gradio as gr
from smolagents import CodeAgent, HfApiModel, DuckDuckGoSearchTool, tool
from smolagents.agents import ActionStep
from time import sleep
import helium
from helium import S, Keys
import requests
from bs4 import BeautifulSoup


# Define a function to extract image URLs from DuckDuckGo search
@tool
def search_images_for_presentation(query: str) -> str:
    """
    Searches for image related to the given query on the web (using DuckDuckGo search).
    Args:
        query: The query to search for images or logos.

    Returns:
        list: A list of image URLs.
    """
    helium.start_chrome(headless=True) 
    helium.go_to("https://duckduckgo.com/")
    search_box = helium.find(S("input[type='text']"))
    search_box.write(query + " image")
    search_box.press(Keys.ENTER)

    # Wait for search results to load
    sleep(3)

    # Extract image links from the page
    soup = BeautifulSoup(helium.get_driver().page_source, "html.parser")
    images = [img["src"] for img in soup.find_all("img") if "src" in img.attrs]

    return images[0]  # Return the top 5 image URLs


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


# Function to run the agent and return image links
def run_agent(query: str):
    model = HfApiModel(model_id="https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud/")
    agent = initialize_agent(model)
    image_urls = agent.run(query)
    return image_urls


# Gradio interface function
def gradio_interface(query: str):
    image_urls = run_agent(query)

    # Format output as clickable links
    links = [f'<a href="{url}" target="_blank">{url}</a>' for url in image_urls]
    
    return image_urls


# Set up Gradio interface
gr.Interface(
    fn=gradio_interface,
    inputs="text",
    outputs="text",
    title="Figure, Image & Logo Finder",
    description="Enter a query to search for relevant images, logos, or figures for your presentation.",
).launch()