#!/usr/bin/env python3
import sys
from selenium import webdriver
from selenium.webdriver.common.by import By
import xml.etree.ElementTree as ET

class ScreenshotGenerator:
    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_page_content(self) -> dict:
        script = """
        function processElement(element) {
            // Check if element is actually visible in viewport
            const rect = element.getBoundingClientRect();
            if (!element.offsetParent && element !== document.body) return null;
            if (rect.width === 0 || rect.height === 0) return null;
            
            // Check if element is within viewport bounds
            const viewportWidth = window.innerWidth;
            const viewportHeight = window.innerHeight;
            if (rect.bottom < 0 || rect.top > viewportHeight || 
                rect.right < 0 || rect.left > viewportWidth) {
                return null;
            }
            
            const styles = window.getComputedStyle(element);
            
            const text = Array.from(element.childNodes)
                .filter(node => node.nodeType === Node.TEXT_NODE)
                .map(node => node.textContent.trim())
                .join(' ');
                
            const interactions = [];
            if (element.tagName === 'BUTTON' || element.tagName === 'A' || 
                element.getAttribute('role') === 'button' || 
                element.getAttribute('onclick') || 
                styles.cursor === 'pointer') {
                interactions.push('click');
            }
            
            if (element.tagName === 'INPUT' || element.tagName === 'TEXTAREA' ||
                element.getAttribute('contenteditable') === 'true') {
                interactions.push('text_input');
                const placeholder = element.getAttribute('placeholder');
                if (placeholder) interactions.push('placeholder:' + placeholder);
            }
            
            if (element.tagName === 'SELECT' || 
                element.getAttribute('role') === 'listbox' ||
                element.getAttribute('type') === 'checkbox' ||
                element.getAttribute('type') === 'radio') {
                interactions.push('select');
            }
            
            if (element.scrollHeight > element.clientHeight || 
                element.scrollWidth > element.clientWidth) {
                interactions.push('scroll');
            }
            
            const children = Array.from(element.children)
                .map(child => processElement(child))
                .filter(child => child !== null);
                
            return {
                tag: element.tagName.toLowerCase(),
                text,
                location: {x: Math.round(rect.left), y: Math.round(rect.top)},
                size: {width: Math.round(rect.width), height: Math.round(rect.height)},
                interactions,
                children
            };
        }
        
        const body = document.body;
        return {
            viewport: {
                scroll_x: window.pageXOffset,
                scroll_y: window.pageYOffset,
                width: document.documentElement.clientWidth,
                height: document.documentElement.clientHeight,
                max_scroll_y: Math.max(document.documentElement.scrollHeight - window.innerHeight, 0),
                min_scroll_x: Math.max(document.documentElement.scrollLeft, 0)
            },
            content: body ? processElement(body) : null
        };
        """
        return self.driver.execute_script(script)

    def format_xml(self, xml_str: str) -> str:
        root = ET.fromstring(xml_str)
        ET.indent(ET.ElementTree(root), space="  ", level=0)
        return ET.tostring(root, encoding='unicode', method='xml')
        
    def generate_xml(self) -> str:
        page_data = self.get_page_content()
        
        viewport = ET.Element("viewport")
        for k, v in page_data['viewport'].items():
            viewport.set(k, str(v))
        
        def process_node(node, parent):
            if node is None:
                return
                
            if node['text']:
                if len(parent) == 0:
                    parent.text = (parent.text or '') + node['text'] + '\n'
                else:
                    parent[-1].tail = (parent[-1].tail or '') + node['text'] + '\n'
                    
            if node['interactions']:
                interactive = ET.SubElement(parent, "interactive")
                interactive.set("x", str(node['location']['x']))
                interactive.set("y", str(node['location']['y']))
                interactive.set("width", str(node['size']['width']))
                interactive.set("height", str(node['size']['height']))
                interactive.set("interactions", ",".join(
                    i for i in node['interactions'] if not i.startswith('placeholder:')
                ))
                
                placeholder = next(
                    (i.split(':', 1)[1] for i in node['interactions'] 
                     if i.startswith('placeholder:')),
                    None
                )
                if placeholder:
                    interactive.set('placeholder', placeholder)
                    
                if node['text']:
                    interactive.text = node['text']
                    
            for child in node['children']:
                process_node(child, parent)
        
        process_node(page_data['content'], viewport)
        return self.format_xml(ET.tostring(viewport, encoding='unicode'))

def main(debug_port: int):
    try:
        generator = ScreenshotGenerator(debug_port)
        xml_output = generator.generate_xml()
        print(xml_output)
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)

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