agents/agent_1_design_inspector.py

"""
Agent 1: Design Inspector
Captures screenshots AND extracts element tree from Figma design file.
Provides dev-mode style specifications for each element.
"""

from typing import Dict, Any
import sys
import os

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from utils.figma_client import FigmaClient
from utils.figma_element_extractor import FigmaElementExtractor, extract_figma_elements


def agent_1_node(state: Dict[str, Any]) -> Dict[str, Any]:
    """
    Capture screenshots AND extract elements from Figma design.
    
    This agent:
    1. Connects to Figma API
    2. Finds Desktop and Mobile frames
    3. Exports screenshots at proper scale
    4. Extracts ALL elements with dev-mode specs
    5. Stores paths, dimensions, and elements in state
    """
    print("\n" + "="*60)
    print("🎨 Agent 1: Design Inspector - Figma Analysis")
    print("="*60)
    
    figma_key = state.get("figma_access_token", "")
    figma_file = state.get("figma_file_key", "")
    execution_id = state.get("execution_id", "")
    logs = state.get("logs", [])
    
    try:
        # Initialize Figma client
        client = FigmaClient(figma_key)
        
        print("\n  📥 Step 1: Fetching Figma file structure...")
        
        # Get full file data for element extraction
        file_data = client.get_file(figma_file)
        file_name = file_data.get("name", "Unknown")
        print(f"     File: {file_name}")
        logs.append(f"📄 Figma file: {file_name}")
        
        # Export frames as screenshots
        print("\n  📸 Step 2: Exporting frame screenshots...")
        screenshots, dimensions = client.export_frames_for_comparison(
            file_key=figma_file,
            output_dir="data/figma",
            execution_id=execution_id
        )
        
        if not screenshots:
            raise ValueError("No frames found in Figma file. Ensure frames are named with 'Desktop' or 'Mobile'.")
        
        # Log screenshots
        for viewport, path in screenshots.items():
            dims = dimensions.get(viewport, {})
            print(f"     ✅ {viewport}: {dims.get('width', '?')}x{dims.get('height', '?')}px")
            logs.append(f"📸 Figma {viewport}: {dims.get('width', '?')}x{dims.get('height', '?')}px")
        
        # Extract elements from each frame
        print("\n  🔍 Step 3: Extracting UI elements (dev mode specs)...")
        
        all_elements = {}
        all_summaries = {}
        
        # Get frame info for each viewport
        frames = client.get_frame_nodes(figma_file)
        
        for frame in frames:
            viewport = frame.get("viewport")
            frame_id = frame.get("id")
            frame_name = frame.get("name")
            
            if not viewport or not frame_id:
                continue
            
            print(f"\n     📱 Processing {viewport} frame: {frame_name}")
            
            try:
                elements, summary = extract_figma_elements(file_data, frame_id, viewport)
                
                # Convert to dict for JSON serialization
                all_elements[viewport] = [e.to_dict() for e in elements]
                all_summaries[viewport] = summary
                
                print(f"        Total elements: {summary['total_elements']}")
                print(f"        Interactive: {summary['interactive_elements']}")
                print(f"        Text elements: {summary['text_elements']}")
                
                # Show element type breakdown
                type_counts = summary.get('by_type', {})
                if type_counts:
                    top_types = sorted(type_counts.items(), key=lambda x: -x[1])[:5]
                    types_str = ", ".join([f"{t}:{c}" for t, c in top_types])
                    print(f"        Types: {types_str}")
                
                logs.append(f"🔍 Figma {viewport}: {summary['total_elements']} elements extracted")
                
            except Exception as e:
                print(f"        ⚠️ Element extraction failed: {str(e)}")
                logs.append(f"⚠️ {viewport} element extraction failed")
                all_elements[viewport] = []
                all_summaries[viewport] = {"total_elements": 0, "error": str(e)}
        
        # Summary
        total_elements = sum(s.get('total_elements', 0) for s in all_summaries.values())
        print(f"\n  ✅ Complete: {len(screenshots)} screenshots, {total_elements} elements extracted")
        
        return {
            "figma_screenshots": screenshots,
            "figma_dimensions": dimensions,
            "figma_elements": all_elements,
            "figma_element_summary": all_summaries,
            "status": "figma_captured",
            "logs": logs
        }
        
    except Exception as e:
        error_msg = f"Failed to analyze Figma file: {str(e)}"
        print(f"\n  ❌ {error_msg}")
        logs.append(f"❌ {error_msg}")
        
        import traceback
        traceback.print_exc()
        
        return {
            "status": "figma_capture_failed",
            "error_message": error_msg,
            "logs": logs
        }