#!/usr/bin/env python3
import sys
from selenium import webdriver
from selenium.webdriver.common.by import By
from selenium.webdriver.remote.webelement import WebElement
import xml.etree.ElementTree as ET
from typing import Dict, List, Set, Optional, Union
from dataclasses import dataclass

@dataclass
class DOMNode:
    element: WebElement
    children: List['DOMNode']
    text: str
    tag: str
    location: Dict[str, int]
    size: Dict[str, int]
    interactions: List[str]
    
    def get_content(self) -> List[Union[ET.Element, str]]:
        content = []
        
        # Add text as plain string if present 
        if self.text:
            content.append(self.text)
            
        # Add interactive element if present
        if self.interactions:
            attrs = {
                "x": str(self.location["x"]),
                "y": str(self.location["y"]), 
                "width": str(self.size["width"]),
                "height": str(self.size["height"])
            }
            
            # Split interactions and placeholders
            base_interactions = []
            placeholder = None
            
            for interaction in self.interactions:
                if interaction.startswith("placeholder:"):
                    placeholder = interaction.split(":", 1)[1]
                else:
                    base_interactions.append(interaction)
                    
            attrs["interactions"] = ",".join(base_interactions)
            if placeholder:
                attrs["placeholder"] = placeholder
                
            interactive = ET.Element("interactive", attrs)
            if self.text:
                interactive.text = self.text
            content.append(interactive)
            
        # Add all children's content
        for child in self.children:
            content.extend(child.get_content())
            
        return content

class DOMTree:
    def __init__(self, root_element: WebElement):
        print(f"\nBuilding DOM tree for {root_element.tag_name}")
        self.root = self._build_tree(root_element)
        
    def _get_element_text(self, element: WebElement) -> str:
        script = """
            let element = arguments[0];
            let text = '';
            for (let node of element.childNodes) {
                if (node.nodeType === Node.TEXT_NODE) {
                    text += node.textContent;
                }
            }
            return text.trim();
        """
        text = element.parent.execute_script(script, element)
        if text:
            print(f"Got text: {text[:50]}...")
        return text
        
    def _get_interactions(self, element: WebElement) -> List[str]:
        interactions = []
        
        if (element.tag_name in {"button", "a"} or
            element.get_attribute("role") in {"button", "link"} or
            element.get_attribute("onclick") or
            element.value_of_css_property("cursor") == "pointer"):
            interactions.append("click")
            
        if (element.tag_name in {"input", "textarea"} or
            element.get_attribute("contenteditable") == "true" or
            element.get_attribute("role") == "textbox"):
            interactions.append("text_input")
            # Get placeholder if available
            placeholder = element.get_attribute("placeholder")
            if placeholder:
                interactions.append(f"placeholder:{placeholder}")
            
        if (element.tag_name == "select" or
            element.get_attribute("role") == "listbox" or
            element.get_attribute("type") in {"checkbox", "radio"}):
            interactions.append("select")
            
        script = """
            let el = arguments[0];
            return el.scrollHeight > el.clientHeight ||
                   el.scrollWidth > el.clientWidth;
        """
        if element.parent.execute_script(script, element):
            interactions.append("scroll")
            
        if interactions:
            print(f"Found interactions for {element.tag_name}: {interactions}")
        return interactions
        
    def _build_tree(self, element: WebElement) -> DOMNode:
        print(f"\nProcessing <{element.tag_name}>")
        
        if not element.is_displayed():
            print(f"Skipping invisible {element.tag_name}")
            return None
            
        children = []
        for child in element.find_elements(By.XPATH, "./*"):
            if child_node := self._build_tree(child):
                children.append(child_node)
                
        text = self._get_element_text(element)
        interactions = self._get_interactions(element)
                
        return DOMNode(
            element=element,
            children=children,
            text=text,
            tag=element.tag_name,
            location=element.location,
            size=element.size,
            interactions=interactions
        )

class ScreenshotGenerator:
    def __init__(self, debug_port: int):
        print(f"\nInitializing ScreenshotGenerator with port {debug_port}")
        options = webdriver.ChromeOptions()
        options.add_experimental_option("debuggerAddress", f"127.0.0.1:{debug_port}")
        self.driver = webdriver.Chrome(options=options)
        
    def get_viewport_metrics(self) -> Dict[str, str]:
        metrics = {
            "scroll_x": self.driver.execute_script("return window.pageXOffset;"),
            "scroll_y": self.driver.execute_script("return window.pageYOffset;"),
            "width": self.driver.execute_script("return document.documentElement.clientWidth;"),
            "height": self.driver.execute_script("return document.documentElement.clientHeight;"),
            "max_scroll_y": self.driver.execute_script(
                "return Math.max(document.documentElement.scrollHeight - window.innerHeight, 0);"
            ),
            "min_scroll_x": self.driver.execute_script(
                "return Math.max(document.documentElement.scrollLeft, 0);"
            )
        }
        print(f"\nViewport metrics: {metrics}")
        return {k: str(v) for k, v in metrics.items()}
        
    def generate_xml(self) -> str:
        metrics = self.get_viewport_metrics()
        metrics_str = " ".join(f'{k}="{v}"' for k, v in metrics.items())
        
        output = [f'<viewport {metrics_str}>']
        seen_texts = set()
        last_was_text = False
        
        body = self.driver.find_element(By.TAG_NAME, "body")
        dom_tree = DOMTree(body)
        
        for content in dom_tree.root.get_content():
            if isinstance(content, str):
                if content not in seen_texts:
                    if last_was_text:
                        output.append("\n")
                    output.append(content)
                    seen_texts.add(content)
                    last_was_text = True
            else:
                output.append(ET.tostring(content, encoding="unicode"))
                last_was_text = False
                
        output.append('</viewport>')
        return "".join(output)

def main(debug_port: int):
    generator = ScreenshotGenerator(debug_port)
    print("\nFinal XML output:")
    print(generator.generate_xml())

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: itb_screenshot <debug_port>")
        sys.exit(1)
    main(int(sys.argv[1]))