from io import BytesIO
from time import sleep

import helium
from PIL import Image
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.common.keys import Keys
from smolagents import ActionStep, CodeAgent, DuckDuckGoSearchTool, tool

from model import model  # needs MODEL_PROVIDER=openai: the agent reads screenshots

# Explains helium to the model, so it knows how to drive the browser
helium_instructions = """
You can use helium to access websites. Don't bother about the helium driver, it's already managed.
We've already ran "from helium import *", so you can call these functions directly:

go_to('github.com/trending')              # open a page
click("Top products")                     # click a link or button by its visible text
click(Link("Top products"))               # click a link specifically
scroll_down(num_pixels=1200)              # scroll down roughly one viewport
if Text('Accept cookies?').exists():      # check if an element is on the page
    click('I accept')

Rules:
- Never try to log in to a page.
- Use the close_popups tool to dismiss pop-ups; do not use helium to close them.
- Use the search_item_ctrl_f tool to jump to text on a long page instead of scrolling a lot.
- Use .exists() to check an element before clicking it.
- Stop after each action to look at the screenshot of the result before your next step.
- Keep each step small: one go_to or click per step, then print what you need.
"""

chrome_options = webdriver.ChromeOptions()
chrome_options.add_argument("--force-device-scale-factor=1")
chrome_options.add_argument("--window-size=1000,1350")
chrome_options.add_argument("--disable-pdf-viewer")
chrome_options.add_argument("--window-position=0,0")

# helium opens Chrome and shares the same driver with our selenium-based tools
driver = helium.start_chrome(headless=False, options=chrome_options)


@tool
def search_item_ctrl_f(text: str, nth_result: int = 1) -> str:
    """
    Searches for text on the current page via Ctrl + F and jumps to the nth occurrence.
    Args:
        text: The text to search for
        nth_result: Which occurrence to jump to (default: 1)
    """
    elements = driver.find_elements(By.XPATH, f"//*[contains(text(), '{text}')]")
    if nth_result > len(elements):
        raise Exception(f"Match n°{nth_result} not found (only {len(elements)} matches found)")
    result = f"Found {len(elements)} matches for '{text}'."
    elem = elements[nth_result - 1]
    driver.execute_script("arguments[0].scrollIntoView(true);", elem)
    result += f"Focused on element {nth_result} of {len(elements)}"
    return result


@tool
def go_back() -> None:
    """Goes back to previous page."""
    driver.back()


@tool
def close_popups() -> str:
    """
    Closes any visible modal or pop-up on the page. Use this to dismiss pop-up windows! This does not work on cookie consent banners.
    """
    webdriver.ActionChains(driver).send_keys(Keys.ESCAPE).perform()
    return "Sent Escape to close pop-ups."


# Passed to the agent as a step_callback: runs at the end of every step,
# so the model sees a fresh screenshot of the browser before its next step.
def save_screenshot(memory_step: ActionStep, agent: CodeAgent) -> None:
    sleep(1.0)  # Let JavaScript animations happen before taking the screenshot
    current_driver = helium.get_driver()
    if current_driver is None:
        return
    current_step = memory_step.step_number
    for previous_step in agent.memory.steps:  # Remove older screenshots to keep the context small
        if isinstance(previous_step, ActionStep) and previous_step.step_number <= current_step - 2:
            previous_step.observations_images = None
    png_bytes = current_driver.get_screenshot_as_png()
    image = Image.open(BytesIO(png_bytes))
    print(f"Captured a browser screenshot: {image.size} pixels")
    memory_step.observations_images = [image.copy()]  # Copy so the image persists after this function

    # Update observations with current URL
    url_info = f"Current url: {current_driver.current_url}"
    memory_step.observations = url_info if memory_step.observations is None else memory_step.observations + "\n" + url_info


agent = CodeAgent(
    tools=[DuckDuckGoSearchTool(), go_back, close_popups, search_item_ctrl_f],
    model=model,
    additional_authorized_imports=["helium"],
    step_callbacks=[save_screenshot],
    max_steps=20,
    verbosity_level=2,
)

# Make helium's functions (go_to, click, ...) available inside the agent's code
agent.python_executor("from helium import *")


def main() -> None:
    agent.run("""
I am Alfred, the butler of Wayne Manor, responsible for verifying the identity of guests at party. A superhero has arrived at the entrance claiming to be Wonder Woman, but I need to confirm if she is who she says she is.

Please search for images of Wonder Woman and generate a detailed visual description based on those images. Additionally, navigate to Wikipedia to gather key details about her appearance. With this information, I can determine whether to grant her access to the event.
""" + helium_instructions)


if __name__ == "__main__":
    try:
        main()
    finally:
        helium.kill_browser()  # always close Chrome, even if the run fails