cdp_graph_generator.py

import os
import json

def generate_mermaid_spatial_graph(snapshot_data, output_dir=None):
    """
    Parses raw CDP snapshot data and generates an ultra-lightweight Mermaid.js 
    spatial topology diagram and structured interactive map for LLMs.
    """
    if isinstance(snapshot_data, str):
        with open(snapshot_data, "r", encoding="utf-8") as f:
            data = json.load(f)
    else:
        data = snapshot_data

    strings = data.get("strings", [])
    documents = data.get("documents", [])

    def get_str(idx):
        if idx == -1 or idx is None:
            return ""
        if 0 <= idx < len(strings):
            return strings[idx]
        return ""

    computed_style_keys = ['display', 'visibility', 'opacity', 'pointer-events', 'cursor']
    interactive_elements = []

    for doc_idx, doc in enumerate(documents):
        nodes = doc.get("nodes", {})
        layout = doc.get("layout", {})
        
        node_names = nodes.get("nodeName", [])
        node_types = nodes.get("nodeType", [])
        node_values = nodes.get("nodeValue", [])
        parent_indices = nodes.get("parentIndex", [])
        attributes_list = nodes.get("attributes", [])
        
        # Check for CDP isClickable index list
        is_clickable_data = nodes.get("isClickable", {})
        clickable_indices = set()
        if isinstance(is_clickable_data, dict):
            clickable_indices = set(is_clickable_data.get("index", []))
        elif isinstance(is_clickable_data, list):
            clickable_indices = set(is_clickable_data)

        # Build child mapping
        child_map = {}
        for child_idx, parent_idx in enumerate(parent_indices):
            if parent_idx != -1:
                child_map.setdefault(parent_idx, []).append(child_idx)

        # Map layout properties
        node_indices = layout.get("nodeIndex", [])
        bounds = layout.get("bounds", [])
        styles = layout.get("styles", [])
        
        node_layout_map = {}
        for L, node_idx in enumerate(node_indices):
            node_layout_map[node_idx] = {
                "bounds": bounds[L] if L < len(bounds) else [0, 0, 0, 0],
                "styles": styles[L] if L < len(styles) else []
            }

        # Recursive text extractor
        def extract_text(idx):
            n_type = node_types[idx]
            if n_type == 3:  # Text node
                return get_str(node_values[idx]).strip()
            parts = []
            for child_idx in child_map.get(idx, []):
                child_txt = extract_text(child_idx)
                if child_txt:
                    parts.append(child_txt)
            return " ".join(parts)

        # Helper to get node attributes
        def get_attributes(idx):
            attrs = {}
            if idx < len(attributes_list):
                pairs = attributes_list[idx]
                for i in range(0, len(pairs), 2):
                    if i + 1 < len(pairs):
                        attrs[get_str(pairs[i])] = get_str(pairs[i+1])
            return attrs

        content_width = doc.get("contentWidth", 1920)

        for idx in range(len(node_names)):
            if node_types[idx] != 1:  # Element nodes only
                continue

            tag_name = get_str(node_names[idx]).upper()
            attrs = get_attributes(idx)
            role = attrs.get("role", "").lower()
            tabindex = attrs.get("tabindex", "")

            # 🔍 Layout & Bounds Extraction
            layout_data = node_layout_map.get(idx)
            if not layout_data:
                continue

            x, y, w, h = layout_data["bounds"]

            # 1. Bounds Filter
            if w <= 0 or h <= 0:
                continue
            if x < -20 or y < -20 or x > (content_width + 50):
                continue

            # 🔍 Style Extraction
            style_indices = layout_data["styles"]
            resolved_styles = {}
            for s_idx, key in enumerate(computed_style_keys):
                if s_idx < len(style_indices):
                    resolved_styles[key] = get_str(style_indices[s_idx])

            display = resolved_styles.get("display", "")
            visibility = resolved_styles.get("visibility", "")
            opacity_str = resolved_styles.get("opacity", "1")
            pointer_events = resolved_styles.get("pointer-events", "")
            cursor = resolved_styles.get("cursor", "")

            try:
                opacity = float(opacity_str)
            except ValueError:
                opacity = 1.0

            # 2. Styles Visibility & Pointer Filter
            if display == "none" or visibility == "hidden" or opacity < 0.1:
                continue
            if pointer_events == "none":
                continue

            # 3. Interactive Classifier
            is_clickable = idx in clickable_indices
            is_pointer_cursor = cursor == "pointer"
            is_standard_tag = tag_name in ["BUTTON", "A", "INPUT", "TEXTAREA", "SELECT"]
            is_interactive_role = role in ["button", "link", "checkbox", "radio", "textbox"]
            has_tabindex = bool(tabindex and tabindex != "-1")

            if not (is_clickable or is_pointer_cursor or is_standard_tag or is_interactive_role or has_tabindex):
                continue

            text = extract_text(idx)
            aria_label = attrs.get("aria-label", "")
            placeholder = attrs.get("placeholder", "")
            element_class = attrs.get("class", "").lower()
            
            # Filter out pure vector structural tags unless they have explicit ARIA labels or text
            if tag_name in ["SVG", "PATH", "G", "CIRCLE", "POLYGON", "RECT", "LINE"] and not (aria_label or text):
                continue

            label = text or aria_label or placeholder or attrs.get("value", "") or f"<{tag_name}>"
            label = label.replace("\n", " ").replace('"', "'").strip()
            if len(label) > 60:
                label = label[:57] + "..."

            # Calculate exact center click coordinates
            center_x = round(x + w / 2.0, 1)
            center_y = round(y + h / 2.0, 1)

            # Detect if element belongs to an active Modal / Popup Overlay
            is_modal = False
            modal_keywords = ["send without a note", "add a note", "how do you know", "weekly connection request limit", "send now", "dismiss", "got it"]
            lbl_lower = label.lower()
            if any(k in lbl_lower for k in modal_keywords) or "modal" in role or "dialog" in role or "artdeco-modal" in element_class or attrs.get("aria-modal") == "true":
                is_modal = True

            interactive_elements.append({
                "node_idx": idx,
                "tag": tag_name,
                "label": label,
                "x": round(x, 1),
                "y": round(y, 1),
                "width": round(w, 1),
                "height": round(h, 1),
                "click_x": center_x,
                "click_y": center_y,
                "is_clickable": is_clickable or is_pointer_cursor,
                "has_text": bool(text or aria_label),
                "is_modal": is_modal
            })

    # 🧹 SPATIAL DEDUPLICATION & MODAL PRIORITIZATION
    interactive_elements.sort(key=lambda el: (0 if el.get("is_modal") else 1, el["y"], el["x"]))

    deduped_elements = []
    for el in interactive_elements:
        duplicate_found = False
        for existing in deduped_elements:
            if abs(existing["click_x"] - el["click_x"]) <= 15 and abs(existing["click_y"] - el["click_y"]) <= 15:
                duplicate_found = True
                if not existing["has_text"] and el["has_text"]:
                    existing["label"] = el["label"]
                    existing["tag"] = el["tag"]
                    existing["has_text"] = True
                elif el["tag"] in ["BUTTON", "A", "INPUT"] and existing["tag"] not in ["BUTTON", "A", "INPUT"]:
                    existing["tag"] = el["tag"]
                break
                
        if not duplicate_found:
            deduped_elements.append(el)

    # Separate Modal elements from background elements
    modal_elements = [el for el in deduped_elements if el.get("is_modal")]
    regular_elements = [el for el in deduped_elements if not el.get("is_modal")]
    
    # Calculate Modal Bounding Box Envelope if modal elements exist
    modal_envelope = None
    if modal_elements:
        min_x = min(el["x"] for el in modal_elements) - 20
        min_y = min(el["y"] for el in modal_elements) - 20
        max_x = max(el["x"] + el["width"] for el in modal_elements) + 20
        max_y = max(el["y"] + el["height"] for el in modal_elements) + 20
        if (max_x - min_x) < 300:
            cx = (min_x + max_x) / 2.0
            min_x, max_x = cx - 200, cx + 200
        if (max_y - min_y) < 150:
            cy = (min_y + max_y) / 2.0
            min_y, max_y = cy - 120, cy + 120
        modal_envelope = (max(0, min_x), max(0, min_y), min(1280, max_x), min(800, max_y))

    # Check which background elements fall INSIDE the modal envelope (Covered under Popup)
    covered_elements = []
    uncovered_elements = []
    if modal_envelope:
        m_x1, m_y1, m_x2, m_y2 = modal_envelope
        for el in regular_elements:
            if m_x1 <= el["click_x"] <= m_x2 and m_y1 <= el["click_y"] <= m_y2:
                el["is_covered"] = True
                covered_elements.append(el)
            else:
                el["is_covered"] = False
                uncovered_elements.append(el)
    else:
        uncovered_elements = regular_elements

    # Order elements: Modal Overlay FIRST, then Uncovered Background, then Covered Background
    ordered_elements = modal_elements + uncovered_elements + covered_elements

    # Re-assign clean AI IDs (1, 2, 3...)
    for i, el in enumerate(ordered_elements):
        el["ai_id"] = i + 1

    # 🗺️ BUILD MERMAID.JS SPATIAL DIAGRAM
    mermaid_lines = ["graph TD"]
    mermaid_lines.append('    subgraph Page_Layout["🌐 SPATIAL PAGE TOPOLOGY"]')

    top_region = [el for el in ordered_elements if el["y"] < 200]
    mid_region = [el for el in ordered_elements if 200 <= el["y"] < 600]
    bot_region = [el for el in ordered_elements if el["y"] >= 600]

    def add_region(region_name, elements):
        if not elements:
            return
        mermaid_lines.append(f'        subgraph {region_name}')
        for el in elements:
            badge = "🔴 MODAL BUTTON" if el.get("is_modal") else ("⚠️ COVERED UNDER POPUP" if el.get("is_covered") else ("🔴 CLICKABLE" if el["is_clickable"] else "🟦 INTERACTIVE"))
            node_str = f'            E{el["ai_id"]}["[{el["ai_id"]}] <{el["tag"]}> \'{el["label"]}\'<br/>Center: ({el["click_x"]}, {el["click_y"]}) | {badge}"]'
            mermaid_lines.append(node_str)
        mermaid_lines.append('        end')

    if modal_elements:
        add_region('MODAL_REGION["⚡ ACTIVE MODAL OVERLAY (FOREGROUND)"]', modal_elements)

    add_region('HEADER_REGION["📍 TOP REGION (y < 200px)"]', top_region)
    add_region('MAIN_REGION["📍 MAIN CONTENT REGION (200px <= y < 600px)"]', mid_region)
    add_region('FOOTER_REGION["📍 LOWER REGION (y >= 600px)"]', bot_region)

    mermaid_lines.append('    end')
    mermaid_diagram = "\n".join(mermaid_lines)

    # 📝 BUILD 2D VISUAL ASCII LAYER MAP DIAGRAM FOR NON-VISION LLM PROMPT
    ascii_map = []
    ascii_map.append("========================================================================================")
    ascii_map.append("🌐 2D VISUAL ASCII LAYER MAP (FOREGROUND OVERLAYS + BACKGROUND CONTEXT)")
    ascii_map.append("========================================================================================")
    ascii_map.append("+--------------------------------------------------------------------------------------+")
    
    if modal_elements:
        m_x1, m_y1, m_x2, m_y2 = modal_envelope
        ascii_map.append(f"| 🚨 FOREGROUND ACTIVE POPUP MODAL OVERLAY (BOX AREA: X: {int(m_x1)}-{int(m_x2)}, Y: {int(m_y1)}-{int(m_y2)})")
        ascii_map.append("| +----------------------------------------------------------------------------------+")
        for el in modal_elements:
            ascii_map.append(
                f"| |  [{el['ai_id']}] <{el['tag']}> '{el['label']}' -> Target Center: (X: {el['click_x']}, Y: {el['click_y']}) | 🔴 ACTIVE MODAL BUTTON"
            )
        ascii_map.append("| +----------------------------------------------------------------------------------+")
        ascii_map.append("|")

    if uncovered_elements:
        ascii_map.append("| 🟢 UNCOVERED / VISIBLE BACKGROUND ELEMENTS (SAFE TO CLICK IF NO MODAL ACTION NEEDED):")
        for el in uncovered_elements:
            ascii_map.append(
                f"|   [{el['ai_id']}] <{el['tag']}> '{el['label']}' -> Center: (X: {el['click_x']}, Y: {el['click_y']}) | {'🔴 CLICKABLE' if el['is_clickable'] else '🟦 INTERACTIVE'}"
            )
        ascii_map.append("|")

    if covered_elements:
        ascii_map.append("| ⚠️ COVERED BACKGROUND ELEMENTS (CURRENTLY LOCATED UNDERNEATH THE ACTIVE POPUP MODAL):")
        for el in covered_elements:
            ascii_map.append(
                f"|   [{el['ai_id']}] <{el['tag']}> '{el['label']}' -> Center: (X: {el['click_x']}, Y: {el['click_y']}) | ⚠️ COVERED UNDER POPUP"
            )

    ascii_map.append("+--------------------------------------------------------------------------------------+")
    llm_compact_tree = "\n".join(ascii_map)

    # Output paths
    if not output_dir:
        output_dir = os.path.dirname(os.path.abspath(__file__)) if '__file__' in globals() else "."

    mermaid_file = os.path.join(output_dir, "spatial_graph.mermaid")
    json_file = os.path.join(output_dir, "spatial_graph.json")

    with open(mermaid_file, "w", encoding="utf-8") as f:
        f.write(mermaid_diagram)

    with open(json_file, "w", encoding="utf-8") as f:
        json.dump(deduped_elements, f, indent=2)

    return {
        "elements": deduped_elements,
        "mermaid": mermaid_diagram,
        "llm_compact_tree": llm_compact_tree,
        "mermaid_file": mermaid_file,
        "json_file": json_file
    }

def capture_cdp_and_generate_graph(driver, output_dir=None):
    """
    Captures CDP DOMSnapshot directly from Selenium driver and generates spatial topology tree.
    Stores extracted elements array on driver.current_cdp_elements for coordinate action lookup.
    """
    snapshot = driver.execute_cdp_cmd('DOMSnapshot.captureSnapshot', {
        'computedStyles': ['display', 'visibility', 'opacity', 'pointer-events', 'cursor'],
        'includePaintOrder': True,
        'includeDOMRects': True
    })
    res = generate_mermaid_spatial_graph(snapshot, output_dir)
    driver.current_cdp_elements = res["elements"]
    return res

if __name__ == "__main__":
    current_dir = os.path.dirname(os.path.abspath(__file__))
    cdp_json = os.path.join(current_dir, "cdp_snapshot.json")
    if os.path.exists(cdp_json):
        res = generate_mermaid_spatial_graph(cdp_json, current_dir)
        print("✅ Spatial Graph Generated!")
        print(f"📊 Extracted {len(res['elements'])} visible interactive elements.")
        print("\n--- GENERATED MERMAID.JS GRAPH ---")
        print(res["mermaid"])