#!/usr/bin/env python3
import sys
import os
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)
        
        # Load JavaScript files
        js_dir = os.path.join(os.path.dirname(__file__), "..", "js")
        self.js_code = ""
        for filename in ["viewport.js", "mouse.js", "dom_analyzer.js"]:
            with open(os.path.join(js_dir, filename)) as f:
                self.js_code += f.read() + "\n"

    def get_page_content(self) -> dict:
        return self.driver.execute_script(self.js_code + "; return analyzePageContent();")

    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]))