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

logger = logging.getLogger(__name__)

class InteractElementTool(Tool):
    name = "interact_element"
    description = "Interacts with a web element (click, fill text, press keys)."
    inputs = {
        "selector": {"type": "string", "default": None, "nullable": True, "description": "CSS selector to target element"},
        "text": {"type": "string", "default": None, "nullable": True, "description": "Text to locate element"},
        "action": {"type": "string", "default": "click", "nullable": False, "description": "Action: 'click', 'fill', or 'press'"},
        "input_text": {"type": "string", "default": None, "nullable": True, "description": "Text to input for fill action"},
        "key": {"type": "string", "default": None, "nullable": True, "description": "Key to press (e.g., 'ENTER')"}
    }
    output_type = "string"

    def __init__(self, driver):
        super().__init__()
        self.driver = driver
        self.is_initialized = self.driver is not None  # Check if driver is valid
        logger.debug(f"InteractElementTool initialized: is_initialized={self.is_initialized}")

    def forward(self, selector=None, text=None, action="click", input_text=None, key=None):
        if not self.is_initialized:
            return "Error: InteractElementTool is not initialized"
        try:
            if selector:
                element = self.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)}"