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

def scroll_page(driver, selector=None, num_pixels=1200, direction="down"):
    """
    Scroll the page to a specific element or by a number of pixels.
    
    Args:
        driver: Selenium WebDriver instance
        selector (str): CSS selector to scroll to (optional)
        num_pixels (int): Number of pixels to scroll (default: 1200)
        direction (str): Scroll direction ('down' or 'up') (default: 'down')
    
    Returns:
        str: Result of the scroll action
    """
    try:
        if selector:
            element = driver.find_element(By.CSS_SELECTOR, selector)
            driver.execute_script("arguments[0].scrollIntoView(true);", element)
            return f"Scrolled to element with selector {selector}"
        else:
            if direction == "down":
                scroll_down(num_pixels)
                return f"Scrolled down {num_pixels} pixels"
            elif direction == "up":
                scroll_up(num_pixels)
                return f"Scrolled up {num_pixels} pixels"
            else:
                return f"Invalid direction: {direction}"
    except Exception as e:
        return f"Failed to scroll: {str(e)}"

# Register the tool
scroll_page_tool = Tool(
    name="scroll_page",
    description="Scrolls the page to a specific element or by a number of pixels.",
    inputs={
        "selector": {
            "type": "str",
            "default": None,
            "description": "CSS selector to scroll to"
        },
        "num_pixels": {
            "type": "int",
            "default": 1200,
            "description": "Number of pixels to scroll"
        },
        "direction": {
            "type": "str",
            "default": "down",
            "description": "Scroll direction: 'down' or 'up'"
        }
    },
    output_type="str",
    function=scroll_page
)