utils/dom_element_extractor.py

"""
DOM Element Extractor
Extracts UI elements from websites using Playwright.
Gets computed styles and element properties for comparison with Figma.
"""

import asyncio
from typing import Dict, List, Optional, Any, Tuple
from .element_schema import (
    UIElement, ElementType, ElementBounds, ElementStyles,
    DOM_TYPE_MAP
)


# JavaScript to extract elements and their computed styles
EXTRACTION_SCRIPT = """
() => {
    const results = [];
    
    // Selectors for elements we care about (checkout-focused)
    const selectors = [
        'button',
        'input',
        'textarea',
        'select',
        'a',
        'h1', 'h2', 'h3', 'h4', 'h5', 'h6',
        'p',
        'label',
        'span',
        'img',
        'svg',
        '[role="button"]',
        '[role="link"]',
        '[role="textbox"]',
        '[role="checkbox"]',
        '[role="radio"]',
        // Common checkout classes
        '.btn', '.button',
        '.input', '.field',
        '.price', '.total', '.amount',
        '.card', '.form-group',
        // Text classes (Tailwind)
        '.text-muted-foreground',
        '.text-gray-500', '.text-gray-600', '.text-gray-700',
        // Data attributes
        '[data-testid]',
        '[data-cy]'
    ];
    
    // Get all matching elements
    const elements = new Set();
    selectors.forEach(selector => {
        try {
            document.querySelectorAll(selector).forEach(el => elements.add(el));
        } catch(e) {}
    });
    
    // Also get important containers (forms, main sections)
    document.querySelectorAll('form, main, [role="main"], .checkout, .cart, .payment').forEach(el => elements.add(el));
    
    // Process each element
    let index = 0;
    elements.forEach(el => {
        // Skip hidden elements
        const style = window.getComputedStyle(el);
        if (style.display === 'none' || style.visibility === 'hidden' || style.opacity === '0') {
            return;
        }
        
        // Get bounding rect
        const rect = el.getBoundingClientRect();
        
        // Skip elements outside viewport or too small
        if (rect.width < 5 || rect.height < 5) {
            return;
        }
        
        // Skip elements way off screen
        if (rect.top > 10000 || rect.left > 5000) {
            return;
        }
        
        // Get element info
        const tagName = el.tagName.toLowerCase();
        const id = el.id || '';
        const className = el.className || '';
        const classStr = typeof className === 'string' ? className : '';
        
        // Determine name (prefer aria-label, then id, then class, then tag)
        let name = el.getAttribute('aria-label') || 
                   el.getAttribute('data-testid') ||
                   el.getAttribute('placeholder') ||
                   id ||
                   classStr.split(' ')[0] ||
                   tagName;
        
        // Get text content (direct text only, not nested)
        let textContent = '';
        if (['BUTTON', 'A', 'LABEL', 'H1', 'H2', 'H3', 'H4', 'H5', 'H6', 'P', 'SPAN'].includes(el.tagName)) {
            // Get direct text content
            textContent = Array.from(el.childNodes)
                .filter(node => node.nodeType === Node.TEXT_NODE)
                .map(node => node.textContent.trim())
                .join(' ')
                .trim();
            
            // If no direct text, get innerText but truncate
            if (!textContent && el.innerText) {
                textContent = el.innerText.substring(0, 100).trim();
            }
        }
        
        // Use text content as name if element has meaningful text
        if (textContent && textContent.length > 2 && textContent.length < 80) {
            name = textContent;
        }
        
        // Get placeholder for inputs
        const placeholder = el.getAttribute('placeholder') || '';
        
        // Get input type
        const inputType = el.getAttribute('type') || '';
        
        // IMPORTANT: Get the ACTUAL computed color, not inherited
        // Force a reflow to ensure styles are computed
        el.offsetHeight;
        
        // Re-get computed style after reflow
        const finalStyle = window.getComputedStyle(el);
        
        // Extract computed styles
        const computedStyles = {
            backgroundColor: finalStyle.backgroundColor,
            color: finalStyle.color,
            borderColor: finalStyle.borderColor,
            borderWidth: finalStyle.borderWidth,
            borderRadius: finalStyle.borderRadius,
            borderStyle: finalStyle.borderStyle,
            fontFamily: finalStyle.fontFamily,
            fontSize: finalStyle.fontSize,
            fontWeight: finalStyle.fontWeight,
            lineHeight: finalStyle.lineHeight,
            textAlign: finalStyle.textAlign,
            letterSpacing: finalStyle.letterSpacing,
            paddingTop: finalStyle.paddingTop,
            paddingRight: finalStyle.paddingRight,
            paddingBottom: finalStyle.paddingBottom,
            paddingLeft: finalStyle.paddingLeft,
            opacity: finalStyle.opacity,
            boxShadow: finalStyle.boxShadow
        };
        
        // Debug: Log color for text elements
        if (textContent && textContent.includes('Complete your purchase')) {
            console.log('Found element:', textContent, 'Color:', finalStyle.color);
        }
        
        results.push({
            index: index++,
            tagName: tagName,
            id: id,
            className: classStr,
            name: name,
            textContent: textContent,
            placeholder: placeholder,
            inputType: inputType,
            bounds: {
                x: rect.left + window.scrollX,
                y: rect.top + window.scrollY,
                width: rect.width,
                height: rect.height
            },
            styles: computedStyles,
            isInteractive: ['BUTTON', 'A', 'INPUT', 'TEXTAREA', 'SELECT'].includes(el.tagName) ||
                          el.getAttribute('role') === 'button' ||
                          el.onclick !== null
        });
    });
    
    return results;
}
"""


class DOMElementExtractor:
    """
    Extracts UI elements from a webpage using Playwright.
    Gets computed styles that match what's actually rendered.
    """
    
    def __init__(self):
        self.elements: List[UIElement] = []
    
    async def extract_from_page_async(
        self,
        page,  # Playwright page object
        viewport: str = "desktop"
    ) -> List[UIElement]:
        """
        Extract elements from a Playwright page.
        
        Args:
            page: Playwright page object (already navigated to URL)
            viewport: "desktop" or "mobile"
            
        Returns:
            List of UIElement objects
        """
        self.elements = []
        
        # Execute extraction script
        raw_elements = await page.evaluate(EXTRACTION_SCRIPT)
        
        # Convert to UIElement objects
        for raw in raw_elements:
            element = self._create_element(raw, viewport)
            if element:
                self.elements.append(element)
        
        print(f"     📊 Extracted {len(self.elements)} elements from {viewport} page")
        return self.elements
    
    def _create_element(self, raw: Dict, viewport: str) -> Optional[UIElement]:
        """
        Create a UIElement from raw extracted data.
        """
        # Create bounds
        bounds_data = raw.get("bounds", {})
        bounds = ElementBounds(
            x=bounds_data.get("x", 0),
            y=bounds_data.get("y", 0),
            width=bounds_data.get("width", 0),
            height=bounds_data.get("height", 0)
        )
        
        # Determine element type
        element_type = self._determine_element_type(raw)
        
        # Extract styles
        styles = self._extract_styles(raw.get("styles", {}))
        
        # Create element
        return UIElement(
            id=f"dom_{raw.get('index', 0)}_{raw.get('tagName', 'unknown')}",
            element_type=element_type,
            name=raw.get("name", ""),
            bounds=bounds,
            styles=styles,
            text_content=raw.get("textContent", "") or None,
            placeholder=raw.get("placeholder", "") or None,
            source="website",
            original_type=raw.get("tagName", ""),
            is_interactive=raw.get("isInteractive", False),
            input_type=raw.get("inputType", "") or None
        )
    
    def _determine_element_type(self, raw: Dict) -> ElementType:
        """
        Determine semantic element type from DOM element.
        """
        tag = raw.get("tagName", "").lower()
        class_name = raw.get("className", "").lower()
        text = raw.get("textContent", "").lower()
        input_type = raw.get("inputType", "").lower()
        
        # Check tag mapping first
        if tag in DOM_TYPE_MAP:
            base_type = DOM_TYPE_MAP[tag]
            
            # Refine input types
            if tag == "input":
                if input_type == "checkbox":
                    return ElementType.CHECKBOX
                elif input_type == "radio":
                    return ElementType.RADIO
                elif input_type == "submit" or input_type == "button":
                    return ElementType.BUTTON
                return ElementType.INPUT
            
            return base_type
        
        # Check class names for hints
        if any(kw in class_name for kw in ["btn", "button"]):
            return ElementType.BUTTON
        if any(kw in class_name for kw in ["input", "field", "textbox"]):
            return ElementType.INPUT
        if any(kw in class_name for kw in ["price", "amount", "total", "cost"]):
            return ElementType.PRICE
        if any(kw in class_name for kw in ["card"]):
            return ElementType.CARD
        if any(kw in class_name for kw in ["icon", "ico"]):
            return ElementType.ICON
        if any(kw in class_name for kw in ["badge", "tag", "chip"]):
            return ElementType.BADGE
        
        # Check text content for price detection
        if "$" in text or "€" in text or "£" in text:
            return ElementType.PRICE
        
        return ElementType.UNKNOWN
    
    def _extract_styles(self, raw_styles: Dict) -> ElementStyles:
        """
        Convert computed CSS styles to ElementStyles.
        """
        styles = ElementStyles()
        
        # Background color
        bg = raw_styles.get("backgroundColor", "")
        if bg and bg != "rgba(0, 0, 0, 0)" and bg != "transparent":
            styles.background_color = self._css_color_to_hex(bg)
        
        # Text color
        color = raw_styles.get("color", "")
        if color:
            styles.text_color = self._css_color_to_hex(color)
        
        # Border color
        border_color = raw_styles.get("borderColor", "")
        if border_color:
            styles.border_color = self._css_color_to_hex(border_color)
        
        # Border width (parse "1px" -> 1.0)
        border_width = raw_styles.get("borderWidth", "")
        if border_width:
            styles.border_width = self._parse_px_value(border_width)
        
        # Border radius
        border_radius = raw_styles.get("borderRadius", "")
        if border_radius:
            styles.border_radius = self._parse_px_value(border_radius)
        
        # Border style
        border_style = raw_styles.get("borderStyle", "")
        if border_style and border_style != "none":
            styles.border_style = border_style
        
        # Font family (clean up quotes)
        font_family = raw_styles.get("fontFamily", "")
        if font_family:
            # Take first font in the stack
            styles.font_family = font_family.split(",")[0].strip().strip('"\'')
        
        # Font size
        font_size = raw_styles.get("fontSize", "")
        if font_size:
            styles.font_size = self._parse_px_value(font_size)
        
        # Font weight
        font_weight = raw_styles.get("fontWeight", "")
        if font_weight:
            try:
                styles.font_weight = int(font_weight)
            except ValueError:
                # Handle named weights
                weight_map = {"normal": 400, "bold": 700, "lighter": 300, "bolder": 700}
                styles.font_weight = weight_map.get(font_weight, 400)
        
        # Line height
        line_height = raw_styles.get("lineHeight", "")
        if line_height and line_height != "normal":
            styles.line_height = self._parse_px_value(line_height)
        
        # Text align
        text_align = raw_styles.get("textAlign", "")
        if text_align:
            styles.text_align = text_align
        
        # Letter spacing
        letter_spacing = raw_styles.get("letterSpacing", "")
        if letter_spacing and letter_spacing != "normal":
            styles.letter_spacing = self._parse_px_value(letter_spacing)
        
        # Padding
        styles.padding_top = self._parse_px_value(raw_styles.get("paddingTop", ""))
        styles.padding_right = self._parse_px_value(raw_styles.get("paddingRight", ""))
        styles.padding_bottom = self._parse_px_value(raw_styles.get("paddingBottom", ""))
        styles.padding_left = self._parse_px_value(raw_styles.get("paddingLeft", ""))
        
        # Opacity
        opacity = raw_styles.get("opacity", "")
        if opacity:
            try:
                styles.opacity = float(opacity)
            except ValueError:
                pass
        
        # Box shadow
        box_shadow = raw_styles.get("boxShadow", "")
        if box_shadow and box_shadow != "none":
            styles.box_shadow = box_shadow
        
        return styles
    
    def _css_color_to_hex(self, css_color: str) -> Optional[str]:
        """
        Convert CSS color (rgb, rgba, hex) to hex string.
        """
        if not css_color:
            return None
        
        css_color = css_color.strip()
        
        # Already hex
        if css_color.startswith("#"):
            return css_color.upper()
        
        # rgb() or rgba()
        if css_color.startswith("rgb"):
            try:
                # Extract numbers
                import re
                numbers = re.findall(r'[\d.]+', css_color)
                if len(numbers) >= 3:
                    r, g, b = int(float(numbers[0])), int(float(numbers[1])), int(float(numbers[2]))
                    hex_color = f"#{r:02X}{g:02X}{b:02X}"
                    return hex_color
            except Exception as e:
                print(f"        ⚠️ Color parse error: {css_color} - {e}")
        
        return None
    
    def _parse_px_value(self, value: str) -> Optional[float]:
        """
        Parse CSS pixel value (e.g., "16px") to float.
        """
        if not value:
            return None
        
        try:
            # Remove 'px' and convert
            return float(value.replace("px", "").strip())
        except ValueError:
            return None
    
    def get_interactive_elements(self) -> List[UIElement]:
        """Get only interactive elements."""
        return [e for e in self.elements if e.is_interactive]
    
    def get_elements_by_type(self, element_type: ElementType) -> List[UIElement]:
        """Get elements of a specific type."""
        return [e for e in self.elements if e.element_type == element_type]
    
    def summarize(self) -> Dict[str, Any]:
        """Get a summary of extracted elements."""
        type_counts = {}
        for element in self.elements:
            type_name = element.element_type.value
            type_counts[type_name] = type_counts.get(type_name, 0) + 1
        
        return {
            "total_elements": len(self.elements),
            "interactive_elements": len(self.get_interactive_elements()),
            "by_type": type_counts
        }


async def extract_dom_elements_async(
    url: str,
    viewport_width: int = 1440,
    viewport_height: int = 900,
    viewport_name: str = "desktop"
) -> Tuple[List[UIElement], Dict]:
    """
    Extract elements from a URL.
    
    Args:
        url: Website URL to extract from
        viewport_width: Viewport width in pixels
        viewport_height: Viewport height in pixels
        viewport_name: "desktop" or "mobile"
        
    Returns:
        Tuple of (list of UIElements, summary dict)
    """
    from playwright.async_api import async_playwright
    
    async with async_playwright() as p:
        browser = await p.chromium.launch(headless=True)
        
        try:
            page = await browser.new_page(viewport={"width": viewport_width, "height": viewport_height})
            
            # Set mobile user agent if needed
            if viewport_name == "mobile":
                await page.set_extra_http_headers({
                    "User-Agent": "Mozilla/5.0 (iPhone; CPU iPhone OS 15_0 like Mac OS X) AppleWebKit/605.1.15"
                })
            
            await page.goto(url, wait_until="networkidle", timeout=60000)
            await page.wait_for_timeout(2000)  # Wait for any animations
            
            extractor = DOMElementExtractor()
            elements = await extractor.extract_from_page_async(page, viewport_name)
            summary = extractor.summarize()
            
            return elements, summary
            
        finally:
            await browser.close()


def extract_dom_elements(
    url: str,
    viewport_width: int = 1440,
    viewport_height: int = 900,
    viewport_name: str = "desktop"
) -> Tuple[List[UIElement], Dict]:
    """
    Synchronous wrapper for DOM element extraction.
    """
    return asyncio.run(extract_dom_elements_async(url, viewport_width, viewport_height, viewport_name))