from smolagents.tools import Tool
from helium import write, press, click, Text, Link, S
from selenium.webdriver.common.keys import Keys

def interact_element(driver, selector=None, text=None, action="click", input_text=None, key=None):
    """
    Interact with a web element (click, fill text, press keys).
    
    Args:
        driver: Selenium WebDriver instance
        selector (str): CSS selector to target the element (optional)
        text (str): Text to locate the element (optional, for Helium)
        action (str): Action to perform ('click', 'fill', 'press') (default: 'click')
        input_text (str): Text to input if action='fill' (optional)
        key (str): Key to press if action='press' (e.g., 'ENTER') (optional)
    
    Returns:
        str: Result of the interaction
    """
    try:
        if selector:
            element = driver.find_element(By.CSS_SELECTOR, selector)
        elif text:
            if action == "click" and Link(text).exists():
                element = Link(text).web_element
            else:
                element = Text(text).web_element
        else:
            return "Must provide selector or text"

        if action == "click":
            click(element)
            return f"Clicked element with {selector or text}"
        elif action == "fill":
            if not input_text:
                return "input_text required for fill action"
            write(input_text, into=element)
            return f"Filled {input_text} into element with {selector or text}"
        elif action == "press":
            if not key:
                return "key required for press action"
            key_map = {"ENTER": Keys.ENTER, "TAB": Keys.TAB, "RETURN": Keys.RETURN}
            if key.upper() in key_map:
                element.send_keys(key_map[key.upper()])
            else:
                press(key)
            return f"Pressed {key} on element with {selector or text}"
        else:
            return f"Unknown action: {action}"
    except Exception as e:
        return f"Failed to interact with element {selector or text}: {str(e)}"

# Register the tool
tool = Tool(
    name="interact_element",
    description="Interacts with a web element (click, fill text, press keys).",
    inputs={
        "selector": {"type": "str", "default": None, "description": "CSS selector to target element"},
        "text": {"type": "str", "default": None, "description": "Text to locate element"},
        "action": {"type": "str", "default": "click", "description": "Action: 'click', 'fill', or 'press'"},
        "input_text": {"type": "str", "default": None, "description": "Text to input for fill action"},
        "key": {"type": "str", "default": None, "description": "Key to press (e.g., 'ENTER')"}
    },
    output_type="str",
    function=interact_element
)