utils/figma_client.py

"""
Figma API Client
Handles communication with Figma API to extract design screenshots.
"""

import requests
from typing import Dict, List, Tuple, Optional
from pathlib import Path


class FigmaClient:
    """Client for interacting with Figma API."""
    
    BASE_URL = "https://api.figma.com/v1"
    
    def __init__(self, access_token: str):
        self.access_token = access_token
        self.headers = {
            "X-Figma-Token": access_token
        }
    
    def get_file(self, file_key: str) -> Dict:
        """Get file metadata from Figma."""
        url = f"{self.BASE_URL}/files/{file_key}"
        response = requests.get(url, headers=self.headers)
        response.raise_for_status()
        return response.json()
    
    def get_frame_nodes(self, file_key: str) -> List[Dict]:
        """
        Get all top-level frames from the Figma file.
        Returns frames that match viewport patterns (Desktop/Mobile).
        """
        file_data = self.get_file(file_key)
        frames = []
        
        # Navigate through document -> pages -> frames
        document = file_data.get("document", {})
        for page in document.get("children", []):
            if page.get("type") == "CANVAS":
                for child in page.get("children", []):
                    if child.get("type") == "FRAME":
                        frame_name = child.get("name", "")
                        frame_id = child.get("id", "")
                        bounds = child.get("absoluteBoundingBox", {})
                        
                        # Determine viewport type from name
                        viewport = None
                        if "desktop" in frame_name.lower() or bounds.get("width", 0) >= 1000:
                            viewport = "desktop"
                        elif "mobile" in frame_name.lower() or bounds.get("width", 0) <= 500:
                            viewport = "mobile"
                        
                        if viewport:
                            frames.append({
                                "id": frame_id,
                                "name": frame_name,
                                "viewport": viewport,
                                "width": bounds.get("width", 0),
                                "height": bounds.get("height", 0)
                            })
        
        return frames
    
    def export_frame(
        self, 
        file_key: str, 
        frame_id: str, 
        output_path: str,
        scale: float = 1.0,
        format: str = "png"
    ) -> Tuple[str, Dict[str, int]]:
        """
        Export a frame as an image.
        
        Args:
            file_key: Figma file key
            frame_id: Node ID of the frame to export
            output_path: Where to save the image
            scale: Export scale (1.0 = actual size, 0.5 = half size)
            format: Image format (png, jpg, svg, pdf)
            
        Returns:
            Tuple of (saved_path, dimensions_dict)
        """
        # Get export URL
        url = f"{self.BASE_URL}/images/{file_key}"
        params = {
            "ids": frame_id,
            "scale": scale,
            "format": format
        }
        
        response = requests.get(url, headers=self.headers, params=params)
        response.raise_for_status()
        
        data = response.json()
        image_url = data.get("images", {}).get(frame_id)
        
        if not image_url:
            raise ValueError(f"Could not get export URL for frame {frame_id}")
        
        # Download the image
        img_response = requests.get(image_url)
        img_response.raise_for_status()
        
        # Save to file
        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
        with open(output_path, "wb") as f:
            f.write(img_response.content)
        
        # Get image dimensions
        from PIL import Image
        with Image.open(output_path) as img:
            width, height = img.size
        
        return output_path, {"width": width, "height": height}
    
    def export_frames_for_comparison(
        self,
        file_key: str,
        output_dir: str,
        execution_id: str
    ) -> Tuple[Dict[str, str], Dict[str, Dict[str, int]]]:
        """
        Export all relevant frames for UI comparison.
        
        Automatically finds Desktop and Mobile frames and exports them
        at 1x scale (not 2x) for proper comparison with website screenshots.
        
        Returns:
            Tuple of (screenshot_paths, dimensions)
        """
        frames = self.get_frame_nodes(file_key)
        screenshots = {}
        dimensions = {}
        
        # Group by viewport, prefer larger frames
        viewport_frames = {}
        for frame in frames:
            viewport = frame["viewport"]
            if viewport not in viewport_frames:
                viewport_frames[viewport] = frame
            elif frame["width"] * frame["height"] > viewport_frames[viewport]["width"] * viewport_frames[viewport]["height"]:
                viewport_frames[viewport] = frame
        
        # Export each viewport
        for viewport, frame in viewport_frames.items():
            print(f"  📥 Exporting frame: {frame['name']} ({frame['width']}px width)")
            print(f"     Frame ID: {frame['id']}")
            print(f"     Dimensions: {frame['width']}x{frame['height']}")
            
            output_path = f"{output_dir}/{viewport}_{execution_id}.png"
            
            # Export at scale 1.0 (actual design size)
            # Note: Figma often has designs at 2x, we'll handle normalization in comparison
            saved_path, dims = self.export_frame(
                file_key,
                frame["id"],
                output_path,
                scale=1.0  # Export at 1x to get actual design dimensions
            )
            
            screenshots[viewport] = saved_path
            dimensions[viewport] = dims
            print(f"  ✓ Exported: {saved_path}")
        
        return screenshots, dimensions