screenshot_annotator.py

"""
Screenshot Annotator
Generates annotated screenshots with red circles highlighting visual differences
"""

from typing import List, Dict, Any
from state_schema import VisualDifference
import os


class ScreenshotAnnotator:
    """Annotates screenshots with visual difference indicators."""
    
    @staticmethod
    def annotate_screenshot(
        screenshot_path: str,
        differences: List[VisualDifference],
        output_path: str,
        viewport_name: str = "desktop"
    ) -> bool:
        """
        Annotate screenshot with red circles for differences.
        
        Args:
            screenshot_path: Path to original screenshot
            differences: List of visual differences
            output_path: Path to save annotated screenshot
            viewport_name: Viewport name (desktop/mobile)
        
        Returns:
            True if successful
        """
        try:
            from PIL import Image, ImageDraw
            
            if not os.path.exists(screenshot_path):
                return False
            
            # Load image
            img = Image.open(screenshot_path).convert('RGB')
            draw = ImageDraw.Draw(img, 'RGBA')
            
            # Filter differences for this viewport
            viewport_diffs = [d for d in differences if d.viewport == viewport_name]
            
            # Draw circles and labels for each difference
            circle_radius = 30
            label_offset = 40
            
            for idx, diff in enumerate(viewport_diffs):
                if not diff.location:
                    # Generate random location if not provided
                    x = (idx % 3) * 300 + 150
                    y = (idx // 3) * 300 + 150
                else:
                    x = diff.location.get("x", 150)
                    y = diff.location.get("y", 150)
                
                # Draw circle
                circle_color = ScreenshotAnnotator._get_color_by_severity(diff.severity)
                draw.ellipse(
                    [(x - circle_radius, y - circle_radius),
                     (x + circle_radius, y + circle_radius)],
                    outline=circle_color,
                    width=3
                )
                
                # Draw number label
                label_number = str(idx + 1)
                draw.text(
                    (x - 8, y - 8),
                    label_number,
                    fill=(255, 255, 255),
                    font=None
                )
            
            # Create output directory
            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            
            # Save annotated image
            img.save(output_path)
            return True
        
        except Exception as e:
            print(f"Error annotating screenshot: {str(e)}")
            return False
    
    @staticmethod
    def _get_color_by_severity(severity: str) -> tuple:
        """Get color based on severity level."""
        if severity == "High":
            return (255, 0, 0, 200)  # Red
        elif severity == "Medium":
            return (255, 165, 0, 200)  # Orange
        else:
            return (0, 255, 0, 200)  # Green
    
    @staticmethod
    def create_comparison_image(
        figma_screenshot: str,
        website_screenshot: str,
        figma_annotated: str,
        website_annotated: str,
        output_path: str
    ) -> bool:
        """
        Create side-by-side comparison image.
        
        Args:
            figma_screenshot: Original Figma screenshot
            website_screenshot: Original website screenshot
            figma_annotated: Annotated Figma screenshot
            website_annotated: Annotated website screenshot
            output_path: Path to save comparison
        
        Returns:
            True if successful
        """
        try:
            from PIL import Image
            
            # Load annotated images
            figma_img = Image.open(figma_annotated)
            website_img = Image.open(website_annotated)
            
            # Resize to same height
            max_height = max(figma_img.height, website_img.height)
            figma_img = figma_img.resize(
                (int(figma_img.width * max_height / figma_img.height), max_height)
            )
            website_img = website_img.resize(
                (int(website_img.width * max_height / website_img.height), max_height)
            )
            
            # Create side-by-side image
            total_width = figma_img.width + website_img.width + 20
            comparison = Image.new('RGB', (total_width, max_height), (255, 255, 255))
            
            comparison.paste(figma_img, (0, 0))
            comparison.paste(website_img, (figma_img.width + 20, 0))
            
            # Create output directory
            os.makedirs(os.path.dirname(output_path), exist_ok=True)
            
            # Save comparison
            comparison.save(output_path)
            return True
        
        except Exception as e:
            print(f"Error creating comparison image: {str(e)}")
            return False


def annotate_all_screenshots(state: Dict[str, Any]) -> Dict[str, Any]:
    """
    Annotate all screenshots with visual differences.
    
    Args:
        state: Current workflow state
    
    Returns:
        Updated state with annotated screenshots
    """
    print("\n📸 Generating Annotated Screenshots...")
    
    annotator = ScreenshotAnnotator()
    
    figma_screenshots = state.get("figma_screenshots", {})
    website_screenshots = state.get("website_screenshots", {})
    visual_differences = state.get("visual_differences", [])
    execution_id = state.get("execution_id", "unknown")
    
    # Convert dict differences back to VisualDifference objects for the annotator
    vd_objects = []
    for d in visual_differences:
        if isinstance(d, dict):
            vd_objects.append(VisualDifference(
                category=d.get("category", "visual"),
                severity=d.get("severity", "Medium"),
                issue_id=d.get("issue_id", "unknown"),
                title=d.get("title", "Difference"),
                description=d.get("description", ""),
                design_value=d.get("design_value", ""),
                website_value=d.get("website_value", ""),
                viewport=d.get("viewport", "desktop"),
                location=d.get("location")
            ))
        else:
            vd_objects.append(d)

    annotated_screenshots = state.get("annotated_screenshots", {})
    
    for viewport_name, figma_path in figma_screenshots.items():
        if viewport_name not in website_screenshots:
            continue
        
        print(f"  📝 Annotating {viewport_name} screenshots...")
        
        website_path = website_screenshots[viewport_name]
        
        figma_annotated = f"data/annotated/{execution_id}_figma_{viewport_name}.png"
        website_annotated = f"data/annotated/{execution_id}_website_{viewport_name}.png"
        
        # Annotate Figma screenshot
        if annotator.annotate_screenshot(figma_path, vd_objects, figma_annotated, viewport_name):
            print(f"    ✓ Figma annotated: {figma_annotated}")
        
        # Annotate website screenshot
        if annotator.annotate_screenshot(website_path, vd_objects, website_annotated, viewport_name):
            print(f"    ✓ Website annotated: {website_annotated}")
        
        # Create comparison
        comparison_path = f"data/comparisons/{execution_id}_{viewport_name}_comparison.png"
        if annotator.create_comparison_image(figma_path, website_path, figma_annotated, website_annotated, comparison_path):
            print(f"    ✓ Comparison created: {comparison_path}")
            annotated_screenshots[f"{viewport_name}_comparison"] = comparison_path
    
    state["annotated_screenshots"] = annotated_screenshots
    return state