utils/figma_client.py

"""
Figma Client - Export frames and extract design specifications
Uses Figma REST API to export frames as PNG images
FIXED: Using the correct /images endpoint that preserves frame dimensions
"""

import os
import requests
import logging
from typing import Dict, Any, Optional
import time

logger = logging.getLogger(__name__)


class FigmaClient:
    """Client for interacting with Figma API."""
    
    def __init__(self, access_token: str = None):
        """
        Initialize Figma Client.
        
        Args:
            access_token: Figma API access token
        """
        self.access_token = access_token or os.getenv("FIGMA_ACCESS_TOKEN", "")
        self.base_url = "https://api.figma.com/v1"
        self.headers = {
            "X-Figma-Token": self.access_token
        }
    
    def _make_request(self, method: str, url: str, params: Dict = None, timeout: int = 30) -> requests.Response:
        """
        Make HTTP request with retry logic.
        
        Args:
            method: HTTP method (GET, POST, etc.)
            url: Request URL
            params: Query parameters
            timeout: Request timeout in seconds
        
        Returns:
            Response object
        """
        max_retries = 3
        retry_delay = 1
        
        for attempt in range(max_retries):
            try:
                if method == "GET":
                    response = requests.get(url, headers=self.headers, params=params, timeout=timeout)
                else:
                    raise ValueError(f"Unsupported method: {method}")
                
                response.raise_for_status()
                return response
            
            except requests.exceptions.RequestException as e:
                if attempt < max_retries - 1:
                    print(f"  ⚠️  Request failed: {str(e)}. Retrying in {retry_delay}s...")
                    time.sleep(retry_delay)
                    retry_delay *= 2
                else:
                    raise
    
    def get_file_structure(self, file_key: str) -> Dict[str, Any]:
        """
        Get the file structure from Figma.
        
        Args:
            file_key: Figma file key
        
        Returns:
            File structure data
        """
        url = f"{self.base_url}/files/{file_key}"
        response = self._make_request("GET", url)
        return response.json()
    
    def find_frames(self, file_key: str) -> Dict[str, Dict[str, Any]]:
        """
        Find all top-level frames in the Figma file.
        
        Args:
            file_key: Figma file key
        
        Returns:
            Dictionary mapping frame names to frame data (id, width, height)
        """
        file_data = self.get_file_structure(file_key)
        frames = {}
        
        def traverse_nodes(node, depth=0):
            """Recursively traverse nodes to find frames."""
            # Only look at direct children of canvases (depth 1)
            # This avoids nested components
            if node.get("type") == "FRAME" and depth == 1:
                frames[node.get("name")] = {
                    "id": node.get("id"),
                    "width": node.get("absoluteBoundingBox", {}).get("width"),
                    "height": node.get("absoluteBoundingBox", {}).get("height")
                }
            
            # Only traverse children if we're at canvas level (depth 0)
            if node.get("type") == "CANVAS" and depth == 0:
                if "children" in node:
                    for child in node["children"]:
                        traverse_nodes(child, depth + 1)
        
        # Start traversal from document root
        if "document" in file_data:
            for child in file_data["document"].get("children", []):
                traverse_nodes(child, depth=0)
        
        return frames
    
    def export_frame(
        self,
        file_key: str,
        frame_name: str,
        output_path: str,
        scale: float = 2.0,
        format: str = "png"
    ) -> str:
        """
        Export a frame from Figma as an image.
        Uses the correct /images endpoint that preserves frame dimensions.
        
        Args:
            file_key: Figma file key
            frame_name: Name of the frame to export
            output_path: Path to save the exported image
            scale: Export scale (1.0 = 100%, 2.0 = 200%)
            format: Export format (png, jpg, svg, pdf)
        
        Returns:
            Path to the exported image
        """
        # Get frame ID
        frames = self.find_frames(file_key)
        
        if frame_name not in frames:
            available = list(frames.keys())
            raise ValueError(f"Frame '{frame_name}' not found. Available frames: {available}")
        
        frame_id = frames[frame_name]["id"]
        frame_width = frames[frame_name]["width"]
        frame_height = frames[frame_name]["height"]
        
        print(f"  📥 Exporting frame: {frame_name}")
        print(f"     Frame ID: {frame_id}")
        print(f"     Dimensions: {frame_width}x{frame_height}")
        
        # Use CORRECT endpoint: /images/{file_key} (NOT /files/{file_key}/images)
        # This endpoint returns images with node_id as the key, preserving frame dimensions
        url = f"{self.base_url}/images/{file_key}"
        params = {
            "ids": frame_id,
            "format": format,
            "scale": scale
        }
        
        response = self._make_request("GET", url, params=params)
        image_data = response.json()
        
        # Check for errors
        if image_data.get("error"):
            raise ValueError(f"Figma API error: {image_data.get('error')}")
        
        # Parse the response - images are directly under 'images' key with node_id as key
        images = image_data.get("images", {})
        
        if not images:
            raise ValueError(f"Failed to export frame '{frame_name}' - no image URLs returned")
        
        # The correct endpoint returns the node_id as the key
        if frame_id in images:
            image_url = images[frame_id]
        else:
            # Fallback: use first available image
            image_url = list(images.values())[0]
        
        if not image_url:
            raise ValueError(f"Failed to export frame '{frame_name}' - image URL is empty")
        
        # Download the image
        image_response = requests.get(image_url, timeout=30)
        image_response.raise_for_status()
        
        # Save to file
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        
        with open(output_path, "wb") as f:
            f.write(image_response.content)
        
        logger.info(f"Exported frame '{frame_name}' to {output_path}")
        print(f"  ✓ Exported: {output_path}")
        
        return output_path
    
    def export_frames_by_pattern(
        self,
        file_key: str,
        pattern: str,
        output_dir: str,
        scale: float = 2.0,
        format: str = "png"
    ) -> Dict[str, str]:
        """
        Export multiple frames matching a pattern.
        
        Args:
            file_key: Figma file key
            pattern: Frame name pattern (e.g., "Checkout-*")
            output_dir: Directory to save exported frames
            scale: Export scale
            format: Export format
        
        Returns:
            Dictionary mapping frame names to export paths
        """
        frames = self.find_frames(file_key)
        exported = {}
        
        for frame_name in frames.keys():
            if pattern.replace("*", "") in frame_name:
                output_path = os.path.join(output_dir, f"{frame_name}.{format}")
                try:
                    self.export_frame(file_key, frame_name, output_path, scale, format)
                    exported[frame_name] = output_path
                except Exception as e:
                    logger.error(f"Failed to export frame '{frame_name}': {str(e)}")
        
        return exported
    
    def get_design_specs(self, file_key: str) -> Dict[str, Any]:
        """
        Extract design specifications from Figma file.
        
        Args:
            file_key: Figma file key
        
        Returns:
            Dictionary with design specifications
        """
        file_data = self.get_file_structure(file_key)
        frames = self.find_frames(file_key)
        
        specs = {
            "file_name": file_data.get("name", ""),
            "frames": frames,
            "colors": [],
            "typography": []
        }
        
        return specs


if __name__ == "__main__":
    # Test the client
    import sys
    
    client = FigmaClient()
    
    if len(sys.argv) < 3:
        print("Usage: python figma_client.py <file_key> <frame_name> [output_path]")
        sys.exit(1)
    
    file_key = sys.argv[1]
    frame_name = sys.argv[2]
    output_path = sys.argv[3] if len(sys.argv) > 3 else f"{frame_name}.png"
    
    print(f"Exporting frame: {frame_name}")
    result = client.export_frame(file_key, frame_name, output_path)
    print(f"Saved to: {result}")