#!/usr/bin/env python3
import sys
import os
import argparse
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, debug: bool = False):
        """Initialize the screenshot generator with optional debug mode"""
        self.debug = debug
        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 log(self, msg: str):
        """Log message if verbose mode is enabled"""
        if self.debug:
            print(f"[Debug] {msg}")

    def get_visible_dom(self) -> str:
        """Get only the visible DOM elements with enhanced visibility detection"""
        if self.debug:
            self.log("Starting DOM analysis")
            print("Debug mode enabled, gathering DOM information...")
        
        try:
            script = """
            // Debug helper
            const debug = {
                log: [],
                add: function(msg) {
                    this.log.push(msg);
                }
            };

            function isVisible(element) {
                debug.add(`Checking visibility for ${element.tagName}`);
                
                if (!element || element.nodeType !== Node.ELEMENT_NODE) {
                    debug.add('Not an element node');
                    return false;
                }
                
                const style = window.getComputedStyle(element);
                const rect = element.getBoundingClientRect();
                
                debug.add(`Element: ${element.tagName}, display: ${style.display}, visibility: ${style.visibility}`);
                debug.add(`Dimensions: ${rect.width}x${rect.height}`);
                
                return (
                    style.display !== 'none' &&
                    style.visibility !== 'hidden' &&
                    rect.width > 0 &&
                    rect.height > 0
                );
            }

            function getElementInfo(element) {
                const info = {
                    tag: element.tagName.toLowerCase(),
                    attributes: {},
                    text: element.textContent.trim()
                };
                
                ['id', 'class', 'href', 'src'].forEach(attr => {
                    if (element.hasAttribute(attr)) {
                        info.attributes[attr] = element.getAttribute(attr);
                    }
                });
                
                return info;
            }

            function traverseDOM(element, depth = 0) {
                if (!element) {
                    debug.add('Null element encountered');
                    return '';
                }

                let output = '';
                if (element.nodeType === Node.ELEMENT_NODE && isVisible(element)) {
                    const info = getElementInfo(element);
                    const indent = '  '.repeat(depth);
                    
                    let elementStr = `${indent}${info.tag}`;
                    Object.entries(info.attributes).forEach(([key, value]) => {
                        elementStr += ` ${key}="${value}"`;
                    });
                    
                    if (info.text && !element.children.length) {
                        elementStr += ` [text: ${info.text}]`;
                    }
                    
                    output += elementStr + '\\n';
                    
                    Array.from(element.children).forEach(child => {
                        output += traverseDOM(child, depth + 1);
                    });
                }
                
                return output;
            }

            const dom = traverseDOM(document.documentElement);
            debug.add(`Total output length: ${dom.length}`);

            return {
                result: dom,
                debug: debug.log,
                stats: {
                    url: window.location.href,
                    title: document.title,
                    elementCount: document.getElementsByTagName('*').length,
                    bodyChildCount: document.body ? document.body.children.length : 0
                }
            };
            """
            
            result = self.driver.execute_script(script)
            
            if isinstance(result, dict):
                if self.debug:
                    print("\n=== Page Statistics ===")
                    stats = result.get('stats', {})
                    print(f"Title: {stats.get('title')}")
                    print(f"Element count: {stats.get('elementCount')}")
                    print(f"Body children: {stats.get('bodyChildCount')}")
                    print("\n=== Debug Log ===")
                    for log_entry in result.get('debug', []):
                        print(f"Browser: {log_entry}")
                
                return result.get('result', '')
            return result
            
        except Exception as e:
            if self.debug:
                print(f"Error in get_visible_dom: {str(e)}")
            return ""

    def get_full_dom(self) -> str:
        """Get the full DOM structure"""
        return self.driver.execute_script("""
            function traverseDOM(node, indent = '') {
                if (!node) return '';
                let result = '';
                
                // Add node information
                if (node.nodeType === Node.ELEMENT_NODE) {
                    result += indent + node.tagName.toLowerCase();
                    if (node.id) result += ` id="${node.id}"`;
                    if (node.className) result += ` class="${node.className}"`;
                    result += '\\n';
                    
                    // Process children
                    for (let child of node.children) {
                        result += traverseDOM(child, indent + '  ');
                    }
                }
                return result;
            }
            return traverseDOM(document.documentElement);
        """)

    def generate_xml(self) -> str:
        page_data = self.driver.execute_script(self.js_code + "; return analyzePageContent();")
        
        # Create viewport element with attributes
        viewport = ET.Element("viewport")
        for key, value in page_data['viewport'].items():
            viewport.set(key, str(value))
        
        def process_element(element_data, parent):
            # Add fixed position indication if present
            if element_data.get('fixed'):
                if element_data.get('tag') == 'nav':
                    nav = ET.SubElement(parent, "navigation")
                    nav.set("fixed", "true")
                    nav.set("x", str(element_data['x']))
                    nav.set("y", str(element_data['y']))
                    nav.set("width", str(element_data['width']))
                    nav.set("height", str(element_data['height']))
                    current_parent = nav
                else:
                    fixed_container = ET.SubElement(parent, "fixed")
                    fixed_container.set("x", str(element_data['x']))
                    fixed_container.set("y", str(element_data['y']))
                    fixed_container.set("width", str(element_data['width']))
                    fixed_container.set("height", str(element_data['height']))
                    current_parent = fixed_container
            else:
                current_parent = parent

            # Handle text content
            if element_data.get('text'):
                if len(current_parent) == 0:
                    current_parent.text = (current_parent.text or '') + element_data['text'] + '\n'
                else:
                    current_parent[-1].tail = (current_parent[-1].tail or '') + element_data['text'] + '\n'
            
            # Handle interactive elements
            if element_data.get('interactions'):
                interactive = ET.SubElement(current_parent, "interactive")
                interactive.set("x", str(element_data['x']))
                interactive.set("y", str(element_data['y']))
                interactive.set("width", str(element_data['width']))
                interactive.set("height", str(element_data['height']))
                interactive.set("interactions", ','.join(element_data['interactions']))
                
                if element_data.get('id'):
                    interactive.set("id", element_data['id'])
                elif element_data.get('suggestedClick'):
                    interactive.set("suggested_click", 
                        f"{element_data['suggestedClick']['x']},{element_data['suggestedClick']['y']}")
                
                if element_data.get('placeholder'):
                    interactive.set("placeholder", element_data['placeholder'])
                
                if element_data.get('buttonText'):
                    interactive.text = element_data['buttonText']
            
            # Handle images
            if element_data.get('tag') == 'img':
                img = ET.SubElement(current_parent, "img")
                img.set("x", str(element_data['x']))
                img.set("y", str(element_data['y']))
                img.set("width", str(element_data['width']))
                img.set("height", str(element_data['height']))
                img.set("src", element_data.get('src', ''))
                img.set("alt", element_data.get('alt', ''))
                if element_data.get('id'):
                    img.set("id", element_data['id'])

            # Process children
            for child in element_data.get('children', []):
                process_element(child, current_parent)
        
        # Process all elements
        process_element(page_data['content'], viewport)
        
        return self.format_xml(ET.tostring(viewport, encoding='unicode'))

    def format_xml(self, xml_str: str) -> str:
        """Format XML with proper indentation"""
        root = ET.fromstring(xml_str)
        ET.indent(root, space="  ", level=0)
        return ET.tostring(root, encoding='unicode', method='xml')

def main():
    parser = argparse.ArgumentParser(description='Generate page screenshot in XML format')
    parser.add_argument('debug_port', type=int, help='Chrome debug port')
    parser.add_argument('--debug', action='store_true', help='Enable debug output')
    parser.add_argument('--dom', action='store_true', help='Print full DOM structure')
    parser.add_argument('--visible-dom', action='store_true', help='Print visible DOM structure')
    
    args = parser.parse_args()
    
    try:
        generator = ScreenshotGenerator(args.debug_port, args.debug)
        
        # Basic DOM check when debug is enabled
        if args.debug:
            print("\nPerforming basic DOM check...")
            basic_check = generator.driver.execute_script("""
                return {
                    title: document.title,
                    bodyExists: !!document.body,
                    elementCount: document.getElementsByTagName('*').length,
                    firstElement: document.body ? document.body.firstElementChild.tagName : 'none'
                }
            """)
            print(f"Page title: {basic_check['title']}")
            print(f"Body exists: {basic_check['bodyExists']}")
            print(f"Total elements: {basic_check['elementCount']}")
            print(f"First element: {basic_check['firstElement']}")
        
        if args.dom:
            print("=== Full DOM Structure ===")
            try:
                dom_output = generator.get_full_dom()
                if dom_output.strip():
                    print(dom_output)
                else:
                    print("No DOM content found.")
            except Exception as e:
                print(f"Error retrieving DOM structure: {str(e)}")
        elif args.visible_dom:
            print("=== Visible DOM Structure ===")
            try:
                visible_dom = generator.get_visible_dom()
                if visible_dom.strip():
                    print(visible_dom)
                else:
                    print("No visible DOM content found.")
            except Exception as e:
                print(f"Error retrieving visible DOM structure: {str(e)}")
        else:
            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__":
    main()