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

logger = logging.getLogger(__name__)

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, "nullable": True, "description": "CSS selector to scroll to"},
        "num_pixels": {"type": "integer", "default": 1200, "nullable": False, "description": "Number of pixels to scroll"},
        "direction": {"type": "string", "default": "down", "nullable": False, "description": "Scroll direction: 'down' or 'up'"}
    }
    output_type = "string"

    def __init__(self, driver):
        super().__init__()
        self.driver = driver
        self.is_initialized = self.driver is not None
        logger.debug(f"ScrollPageTool initialized: is_initialized={self.is_initialized}")

    def forward(self, selector=None, num_pixels=1200, direction="down"):
        if not self.is_initialized:
            return "Error: ScrollPageTool is not initialized"
        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)}"