from smolagents.tools import Tool
from helium import scroll_down, S
from selenium.webdriver.common.by import By

def scroll_page(driver, pixels=600, selector=None):
    """
    Scroll the page by a specified number of pixels or to a specific element.
    
    Args:
        driver: Selenium WebDriver instance
        pixels (int): Number of pixels to scroll down (default: 600)
        selector (str): Optional CSS selector to scroll to a specific element
    
    Returns:
        str: Observation message indicating scroll action
    """
    if selector:
        try:
            element = S(selector)
            driver.execute_script("arguments[0].scrollIntoView(true);", element.web_element)
            return f"Scrolled to element with selector: {selector}"
        except Exception as e:
            return f"Failed to scroll to selector {selector}: {str(e)}"
    else:
        driver.execute_script(f"window.scrollBy(0, {pixels});")
        return f"Scrolled down by {pixels} pixels"

# Register the tool
tool = Tool(
    name="scroll_page",
    description="Scrolls the page by a specified number of pixels or to an element with a given CSS selector.",
    inputs={
        "pixels": {"type": "int", "default": 600, "description": "Number of pixels to scroll down"},
        "selector": {"type": "str", "default": None, "description": "CSS selector of the element to scroll to"}
    },
    output_type="str",
    function=scroll_page
)