#!/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
from dataclasses import dataclass

@dataclass
class DOMNode:
    element: WebElement
    children: List['DOMNode']
    own_text: str
    total_text: str
    tag: str
    location: Dict[str, int]
    size: Dict[str, int]
    is_interactive: bool
    attributes: Dict[str, str]

class DOMTree:
    def __init__(self, root_element: WebElement):
        self.root = self._build_tree(root_element)
        
    def _get_element_own_text(self, element: WebElement) -> str:
        script = """
            let element = arguments[0];
            let walker = document.createTreeWalker(
                element,
                NodeFilter.SHOW_TEXT,
                {
                    acceptNode: function(node) {
                        return node.parentElement === element ? 
                            NodeFilter.FILTER_ACCEPT : 
                            NodeFilter.FILTER_REJECT;
                    }
                }
            );
            let text = '';
            let node;
            while (node = walker.nextNode()) {
                text += node.textContent;
            }
            return text.trim();
        """
        return element.parent.execute_script(script, element)
        
    def _is_interactive(self, element: WebElement) -> bool:
        return (
            element.tag_name in {"button", "a", "input", "select", "textarea"} or
            element.get_attribute("role") in {"button", "link", "textbox"} or
            element.get_attribute("onclick") is not None or
            element.get_attribute("contenteditable") == "true"
        )
        
    def _get_attributes(self, element: WebElement) -> Dict[str, str]:
        attrs = {}
        for attr in ['placeholder', 'aria-label', 'alt', 'title', 'role']:
            if value := element.get_attribute(attr):
                attrs[attr] = value
        return attrs
        
    def _build_tree(self, element: WebElement) -> DOMNode:
        children = [
            self._build_tree(child) 
            for child in element.find_elements(By.XPATH, "./*")
        ]
        
        own_text = self._get_element_own_text(element)
        total_text = own_text + " " + " ".join(
            child.total_text for child in children
        )
        total_text = " ".join(total_text.split())
        
        return DOMNode(
            element=element,
            children=children,
            own_text=own_text,
            total_text=total_text,
            tag=element.tag_name,
            location=element.location,
            size=element.size,
            is_interactive=self._is_interactive(element),
            attributes=self._get_attributes(element)
        )

class ScreenshotGenerator:
    IGNORED_TAGS = {'script', 'style', 'meta', 'link', 'noscript'}
    
    def __init__(self, debug_port: int):
        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, int]:
        return {
            "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);"
            )
        }
        
    def should_process_node(self, node: DOMNode) -> bool:
        # Skip invisible elements
        if not node.element.is_displayed():
            return False
            
        # Skip script and style elements
        if node.tag in self.IGNORED_TAGS:
            return False

        return True
        
    def create_xml_element(self, node: DOMNode, parent_element: Optional[ET.Element] = None) -> Optional[ET.Element]:
        attrs = {
            "x": str(node.location["x"]),
            "y": str(node.location["y"]),
            "width": str(node.size["width"]),
            "height": str(node.size["height"])
        }
        
        if node.is_interactive:
            element = ET.Element("interactive", attrs)
            child = ET.SubElement(element, node.tag)
            for attr, value in node.attributes.items():
                child.set(attr, value)
            if node.own_text:
                child.text = node.own_text
            return element
            
        element = ET.Element(node.tag, attrs)
        if node.own_text:
            element.text = node.own_text
        return element
        
    def process_tree(self, node: DOMNode, parent_element: Optional[ET.Element] = None) -> List[ET.Element]:
        if not self.should_process_node(node):
            return []
            
        current_element = self.create_xml_element(node, parent_element)
        if not current_element:
            return []
            
        for child in node.children:
            child_elements = self.process_tree(child, current_element)
            for child_element in child_elements:
                current_element.append(child_element)
                
        return [current_element]
        
    def generate_output(self) -> ET.Element:
        metrics = self.get_viewport_metrics()
        viewport = ET.Element("viewport", {
            "scroll_x": str(metrics["scroll_x"]),
            "scroll_y": str(metrics["scroll_y"]),
            "width": str(metrics["width"]),
            "height": str(metrics["height"]),
            "max_scroll_y": str(metrics["max_scroll_y"])
        })
        
        dom_tree = DOMTree(self.driver.find_element(By.TAG_NAME, "body"))
        for element in self.process_tree(dom_tree.root):
            viewport.append(element)
            
        return viewport

def main(debug_port: int):
    generator = ScreenshotGenerator(debug_port)
    output = generator.generate_output()
    print(ET.tostring(output, encoding="unicode", method="xml"))

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