utils/figma_element_extractor.py

"""
Figma Element Extractor
Extracts all UI elements from Figma with their dev-mode specifications.
Traverses the node tree recursively and extracts properties for comparison.
"""

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


class FigmaElementExtractor:
    """
    Extracts UI elements from Figma file data.
    Mimics Figma Dev Mode by extracting all design specifications.
    """
    
    # Keywords to detect element types from names
    BUTTON_KEYWORDS = ["button", "btn", "cta", "submit", "pay", "checkout", "continue", "add", "remove"]
    INPUT_KEYWORDS = ["input", "field", "textfield", "textarea", "email", "password", "card", "cvv", "expiry"]
    CHECKBOX_KEYWORDS = ["checkbox", "check", "toggle"]
    PRICE_KEYWORDS = ["price", "total", "subtotal", "amount", "$", "cost"]
    HEADING_KEYWORDS = ["heading", "title", "header", "h1", "h2", "h3"]
    
    def __init__(self, frame_bounds: Optional[Dict] = None):
        """
        Initialize extractor.
        
        Args:
            frame_bounds: The bounds of the parent frame (for relative positioning)
        """
        self.frame_bounds = frame_bounds or {"x": 0, "y": 0}
        self.elements: List[UIElement] = []
        self.element_count = 0
    
    def extract_from_frame(self, frame_node: Dict, viewport: str = "desktop") -> List[UIElement]:
        """
        Extract all elements from a Figma frame.
        
        Args:
            frame_node: The frame node from Figma API
            viewport: "desktop" or "mobile"
            
        Returns:
            List of UIElement objects
        """
        self.elements = []
        self.element_count = 0
        
        # Get frame bounds for relative positioning
        self.frame_bounds = frame_node.get("absoluteBoundingBox", {"x": 0, "y": 0})
        
        # Start recursive extraction
        self._extract_node(frame_node, parent_id=None, depth=0)
        
        print(f"     📊 Extracted {len(self.elements)} elements from {viewport} frame")
        return self.elements
    
    def _extract_node(self, node: Dict, parent_id: Optional[str], depth: int) -> Optional[str]:
        """
        Recursively extract a node and its children.
        
        Returns:
            The ID of the created element, or None if skipped
        """
        node_type = node.get("type", "")
        node_name = node.get("name", "")
        node_id = node.get("id", "")
        
        # Skip invisible nodes
        if not node.get("visible", True):
            return None
        
        # Create element
        element = self._create_element(node, parent_id, depth)
        
        if element:
            self.elements.append(element)
            current_id = element.id
            
            # Process children
            children = node.get("children", [])
            child_ids = []
            
            for child in children:
                child_id = self._extract_node(child, current_id, depth + 1)
                if child_id:
                    child_ids.append(child_id)
            
            # Update children IDs
            element.children_ids = child_ids
            
            return current_id
        
        # If we skipped this node, still process children with same parent
        children = node.get("children", [])
        for child in children:
            self._extract_node(child, parent_id, depth)
        
        return None
    
    def _create_element(self, node: Dict, parent_id: Optional[str], depth: int) -> Optional[UIElement]:
        """
        Create a UIElement from a Figma node.
        """
        node_type = node.get("type", "")
        node_name = node.get("name", "")
        node_id = node.get("id", "")
        
        # Get bounds
        bounds_data = node.get("absoluteBoundingBox")
        if not bounds_data:
            return None
        
        # Calculate relative position (relative to frame)
        bounds = ElementBounds(
            x=bounds_data.get("x", 0) - self.frame_bounds.get("x", 0),
            y=bounds_data.get("y", 0) - self.frame_bounds.get("y", 0),
            width=bounds_data.get("width", 0),
            height=bounds_data.get("height", 0)
        )
        
        # Skip very small elements (likely decorative)
        if bounds.width < 5 or bounds.height < 5:
            return None
        
        # Determine element type
        element_type = self._determine_element_type(node, node_name)
        
        # Extract styles
        styles = self._extract_styles(node)
        
        # Extract text content
        text_content = None
        if node_type == "TEXT":
            text_content = node.get("characters", "")
        
        # Check if interactive
        is_interactive = element_type in [
            ElementType.BUTTON, ElementType.INPUT, ElementType.LINK,
            ElementType.CHECKBOX, ElementType.RADIO, ElementType.SELECT
        ]
        
        self.element_count += 1
        
        return UIElement(
            id=f"figma_{node_id}",
            element_type=element_type,
            name=node_name,
            bounds=bounds,
            styles=styles,
            text_content=text_content,
            parent_id=parent_id,
            depth=depth,
            source="figma",
            original_type=node_type,
            is_interactive=is_interactive
        )
    
    def _determine_element_type(self, node: Dict, name: str) -> ElementType:
        """
        Determine the semantic element type from Figma node.
        Uses node type + name heuristics.
        """
        node_type = node.get("type", "")
        name_lower = name.lower()
        
        # Check by Figma node type first
        if node_type == "TEXT":
            # Check if it's a heading or price
            font_size = node.get("style", {}).get("fontSize", 14)
            
            if any(kw in name_lower for kw in self.PRICE_KEYWORDS):
                return ElementType.PRICE
            if any(kw in name_lower for kw in self.HEADING_KEYWORDS) or font_size >= 24:
                return ElementType.HEADING
            if "label" in name_lower:
                return ElementType.LABEL
            return ElementType.TEXT
        
        # Check by name keywords
        if any(kw in name_lower for kw in self.BUTTON_KEYWORDS):
            return ElementType.BUTTON
        
        if any(kw in name_lower for kw in self.INPUT_KEYWORDS):
            return ElementType.INPUT
        
        if any(kw in name_lower for kw in self.CHECKBOX_KEYWORDS):
            return ElementType.CHECKBOX
        
        if any(kw in name_lower for kw in self.PRICE_KEYWORDS):
            return ElementType.PRICE
        
        # Check for image
        if node_type == "IMAGE" or "image" in name_lower or "img" in name_lower:
            return ElementType.IMAGE
        
        # Check for icon (small vectors or frames with icon in name)
        bounds = node.get("absoluteBoundingBox", {})
        is_small = bounds.get("width", 0) < 50 and bounds.get("height", 0) < 50
        if node_type == "VECTOR" or (is_small and ("icon" in name_lower or "ico" in name_lower)):
            return ElementType.ICON
        
        # Check for divider/line
        if node_type == "LINE" or "divider" in name_lower or "separator" in name_lower:
            return ElementType.DIVIDER
        
        # Check for card
        if "card" in name_lower:
            return ElementType.CARD
        
        # Check for link
        if "link" in name_lower:
            return ElementType.LINK
        
        # Default to container for frames/groups
        if node_type in ["FRAME", "GROUP", "COMPONENT", "INSTANCE"]:
            return ElementType.CONTAINER
        
        # Use mapping for other types
        return FIGMA_TYPE_MAP.get(node_type, ElementType.UNKNOWN)
    
    def _extract_styles(self, node: Dict) -> ElementStyles:
        """
        Extract visual styles from a Figma node.
        """
        styles = ElementStyles()
        
        # Background color (from fills)
        fills = node.get("fills", [])
        if fills and len(fills) > 0:
            fill = fills[0]
            if fill.get("type") == "SOLID" and fill.get("visible", True):
                color = fill.get("color", {})
                styles.background_color = rgb_to_hex(
                    color.get("r", 0),
                    color.get("g", 0),
                    color.get("b", 0)
                )
                styles.opacity = fill.get("opacity", 1.0)
        
        # Border/stroke
        strokes = node.get("strokes", [])
        if strokes and len(strokes) > 0:
            stroke = strokes[0]
            if stroke.get("type") == "SOLID" and stroke.get("visible", True):
                color = stroke.get("color", {})
                styles.border_color = rgb_to_hex(
                    color.get("r", 0),
                    color.get("g", 0),
                    color.get("b", 0)
                )
        
        # Border width
        stroke_weight = node.get("strokeWeight")
        if stroke_weight:
            styles.border_width = stroke_weight
        
        # Border radius
        corner_radius = node.get("cornerRadius")
        if corner_radius:
            styles.border_radius = corner_radius
        
        # Typography (for TEXT nodes)
        style_data = node.get("style", {})
        if style_data:
            styles.font_family = style_data.get("fontFamily")
            styles.font_size = style_data.get("fontSize")
            styles.font_weight = style_data.get("fontWeight")
            styles.line_height = style_data.get("lineHeightPx")
            styles.text_align = style_data.get("textAlignHorizontal", "").lower()
            styles.letter_spacing = style_data.get("letterSpacing")
        
        # Text color (from fills for TEXT nodes)
        if node.get("type") == "TEXT" and fills:
            fill = fills[0]
            if fill.get("type") == "SOLID":
                color = fill.get("color", {})
                styles.text_color = rgb_to_hex(
                    color.get("r", 0),
                    color.get("g", 0),
                    color.get("b", 0)
                )
        
        # Padding (from auto-layout)
        padding_left = node.get("paddingLeft")
        padding_right = node.get("paddingRight")
        padding_top = node.get("paddingTop")
        padding_bottom = node.get("paddingBottom")
        
        if padding_left is not None:
            styles.padding_left = padding_left
        if padding_right is not None:
            styles.padding_right = padding_right
        if padding_top is not None:
            styles.padding_top = padding_top
        if padding_bottom is not None:
            styles.padding_bottom = padding_bottom
        
        # Effects (shadows)
        effects = node.get("effects", [])
        for effect in effects:
            if effect.get("type") == "DROP_SHADOW" and effect.get("visible", True):
                shadow = effect
                color = shadow.get("color", {})
                offset = shadow.get("offset", {})
                radius = shadow.get("radius", 0)
                
                # Format as CSS-like box-shadow
                rgba = f"rgba({int(color.get('r', 0)*255)},{int(color.get('g', 0)*255)},{int(color.get('b', 0)*255)},{color.get('a', 1):.2f})"
                styles.box_shadow = f"{offset.get('x', 0)}px {offset.get('y', 0)}px {radius}px {rgba}"
                break
        
        return styles
    
    def get_interactive_elements(self) -> List[UIElement]:
        """Get only interactive elements (buttons, inputs, etc.)."""
        return [e for e in self.elements if e.is_interactive]
    
    def get_text_elements(self) -> List[UIElement]:
        """Get only text elements."""
        return [e for e in self.elements if e.element_type in [
            ElementType.TEXT, ElementType.HEADING, ElementType.LABEL, ElementType.PRICE
        ]]
    
    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()),
            "text_elements": len(self.get_text_elements()),
            "by_type": type_counts,
            "max_depth": max((e.depth for e in self.elements), default=0)
        }


def extract_figma_elements(file_data: Dict, frame_id: str, viewport: str = "desktop") -> Tuple[List[UIElement], Dict]:
    """
    Convenience function to extract elements from Figma file data.
    
    Args:
        file_data: Complete Figma file data from API
        frame_id: ID of the frame to extract from
        viewport: "desktop" or "mobile"
        
    Returns:
        Tuple of (list of UIElements, summary dict)
    """
    # Find the frame in the document
    frame_node = _find_node_by_id(file_data.get("document", {}), frame_id)
    
    if not frame_node:
        raise ValueError(f"Frame with ID {frame_id} not found in Figma file")
    
    extractor = FigmaElementExtractor()
    elements = extractor.extract_from_frame(frame_node, viewport)
    summary = extractor.summarize()
    
    return elements, summary


def _find_node_by_id(node: Dict, target_id: str) -> Optional[Dict]:
    """Recursively find a node by ID in the Figma document tree."""
    if node.get("id") == target_id:
        return node
    
    for child in node.get("children", []):
        result = _find_node_by_id(child, target_id)
        if result:
            return result
    
    return None