import os
import gradio as gr
from smolagents import CodeAgent, LiteLLMModel, tool
from smolagents.agents import ActionStep
import helium
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.common.keys import Keys
from io import BytesIO
from PIL import Image
from datetime import datetime
from dotenv import load_dotenv
from huggingface_hub import login
import tempfile
import logging

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Load environment variables
load_dotenv()
hf_token = os.getenv("HF_TOKEN")
gemini_api_key = os.getenv("GEMINI_API_KEY")

if not hf_token:
    raise ValueError("HF_TOKEN environment variable not set.")
if not gemini_api_key:
    raise ValueError("GEMINI_API_KEY environment variable not set.")

login(hf_token, add_to_git_credential=False)

# Define tools
@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()

# Set ChromeDriver path
os.environ['webdriver.chrome.driver'] = '/usr/local/bin/chromedriver'

# Initialize Chrome driver with error handling
try:
    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("--no-sandbox")
    chrome_options.add_argument("--disable-dev-shm-usage")
    chrome_options.add_argument("--window-position=0,0")
    chrome_options.add_argument("--headless=new")
    driver = helium.start_chrome(headless=True, options=chrome_options)
    logger.info("Chrome driver initialized successfully.")
except Exception as e:
    logger.error(f"Failed to initialize Chrome driver: {str(e)}")
    raise

# Screenshot callback
def save_screenshot(memory_step: ActionStep, agent: CodeAgent) -> Image.Image:
    from time import sleep
    sleep(1.0)
    driver = helium.get_driver()
    current_step = memory_step.step_number
    if driver is not None:
        for previous_memory_step in agent.memory.steps:
            if isinstance(previous_memory_step, ActionStep) and previous_memory_step.step_number <= current_step - 2:
                previous_memory_step.observations_images = None
        png_bytes = driver.get_screenshot_as_png()
        image = Image.open(BytesIO(png_bytes))
        screenshot_dir = tempfile.gettempdir()
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        screenshot_path = f"{screenshot_dir}/screenshot_step_{current_step}_{timestamp}.png"
        image.save(screenshot_path)
        logger.info(f"Saved screenshot to: {screenshot_path}")
        url_info = f"Current url: {driver.current_url}"
        memory_step.observations = (
            url_info if memory_step.observations is None else memory_step.observations + "\n" + url_info
        )
        return image

# Initialize model and agent
model = LiteLLMModel("gemini/gemini-2.0-flash")
agent = CodeAgent(
    tools=[go_back, close_popups, search_item_ctrl_f],
    model=model,
    additional_authorized_imports=["helium"],
    step_callbacks=[save_screenshot],
    max_steps=20,
    verbosity_level=2,
)
agent.python_executor("from helium import *")

# Helium instructions
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 *"
Then you can go to pages!
Code:
go_to('github.com/trending')
```<end_code>

You can directly click clickable elements by inputting the text that appears on them.
Code:
click("Top products")
```<end_code>

If it's a link:
Code:
click(Link("Top products"))
```<end_code>

If you try to interact with an element and it's not found, you'll get a LookupError.
In general stop your action after each button click to see what happens on your screenshot.
Never try to login in a page.

To scroll up or down, use scroll_down or scroll_up with as an argument the number of pixels to scroll from.
Code:
scroll_down(num_pixels=1200) # This will scroll one viewport down
```<end_code>

When you have pop-ups with a cross icon to close, don't try to click the close icon by finding its element or targeting an 'X' element (this most often fails).
Just use your built-in tool `close_popups` to close them:
Code:
close_popups()
```<end_code>

You can use .exists() to check for the existence of an element. For example:
Code:
if Text('Accept cookies?').exists():
    click('I accept')
```<end_code>
"""

# Gradio interface function
def run_agent(url: str, request: str):
    try:
        search_request = f"Please go to {url}. {request}"
        agent_output = agent.run(search_request + helium_instructions)
        screenshot_path = next(
            (f"{tempfile.gettempdir()}/screenshot_step_{step.step_number}_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png"
             for step in agent.memory.steps if isinstance(step, ActionStep) and step.observations_images),
            None
        )
        return agent_output, screenshot_path
    except Exception as e:
        logger.error(f"Agent execution failed: {str(e)}")
        return f"Error: {str(e)}", None
    finally:
        try:
            driver.quit()
            logger.info("Chrome driver closed.")
        except:
            logger.warning("Failed to close Chrome driver.")

# Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Web Navigation Agent")
    url_input = gr.Textbox(label="Enter URL", placeholder="https://example.com")
    request_input = gr.Textbox(label="Enter Request", placeholder="Describe what to do on the website")
    submit_button = gr.Button("Run Agent")
    output_text = gr.Textbox(label="Agent Output")
    output_image = gr.Image(label="Screenshot")
    submit_button.click(
        fn=run_agent,
        inputs=[url_input, request_input],
        outputs=[output_text, output_image]
    )

if __name__ == "__main__":
    demo.launch()