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

class ScrollPageTool(Tool):
    name = "scroll_page"
    description = "Scrolls the page to a specific element or by a number of pixels."
    inputs = {
        "selector": {"type": "string", "default": None, "description": "CSS selector to scroll to"},
        "num_pixels": {"type": "integer", "default": 1200, "description": "Number of pixels to scroll"},
        "direction": {"type": "string", "default": "down", "description": "Scroll direction: 'down' or 'up'"}
    }
    output_type = "string"

    def __init__(self, driver):
        self.driver = driver

    def forward(self, selector=None, num_pixels=1200, direction="down"):
        try:
            if selector:
                element = self.driver.find_element(By.CSS_SELECTOR, selector)
                self.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)}"