utils/image_differ.py

"""
Image Differ Utility
Compares Figma and Website screenshots, generates visual diff overlays.

Key Features:
- Normalizes image sizes (handles Figma 2x export)
- Creates side-by-side comparison
- Highlights differences in red
- Calculates similarity score
"""

import numpy as np
from PIL import Image, ImageDraw, ImageFont
from pathlib import Path
from typing import Dict, List, Tuple, Any
import cv2


class ImageDiffer:
    """
    Compares two images and generates visual diff output.
    """
    
    def __init__(self, output_dir: str = "data/comparisons"):
        self.output_dir = output_dir
        Path(output_dir).mkdir(parents=True, exist_ok=True)
    
    def normalize_images(
        self,
        figma_path: str,
        website_path: str,
        figma_dims: Dict[str, int],
        website_dims: Dict[str, int]
    ) -> Tuple[np.ndarray, np.ndarray, float]:
        """
        Normalize images to the same size for comparison.
        
        Handles Figma 2x export by detecting and adjusting.
        
        Returns:
            Tuple of (figma_array, website_array, scale_factor)
        """
        figma_img = Image.open(figma_path).convert("RGB")
        website_img = Image.open(website_path).convert("RGB")
        
        figma_w, figma_h = figma_img.size
        website_w, website_h = website_img.size
        
        # Detect if Figma is at 2x (common for retina exports)
        # If Figma is roughly 2x the website width, scale it down
        width_ratio = figma_w / website_w
        scale_factor = 1.0
        
        if 1.8 <= width_ratio <= 2.2:
            # Figma is at 2x, scale down to match website
            scale_factor = 0.5
            new_figma_w = int(figma_w * scale_factor)
            new_figma_h = int(figma_h * scale_factor)
            figma_img = figma_img.resize((new_figma_w, new_figma_h), Image.Resampling.LANCZOS)
            print(f"     📐 Detected Figma 2x export, scaled to {new_figma_w}x{new_figma_h}")
        
        # Now resize both to match (use the smaller dimensions)
        target_w = min(figma_img.size[0], website_img.size[0])
        target_h = min(figma_img.size[1], website_img.size[1])
        
        figma_img = figma_img.resize((target_w, target_h), Image.Resampling.LANCZOS)
        website_img = website_img.resize((target_w, target_h), Image.Resampling.LANCZOS)
        
        return np.array(figma_img), np.array(website_img), scale_factor
    
    def calculate_similarity(
        self,
        img1: np.ndarray,
        img2: np.ndarray
    ) -> Tuple[float, np.ndarray]:
        """
        Calculate similarity score between two images.
        
        Uses structural similarity (SSIM) for perceptual comparison.
        
        Returns:
            Tuple of (similarity_score_0_to_100, diff_mask)
        """
        from skimage.metrics import structural_similarity as ssim
        
        # Convert to grayscale for SSIM
        gray1 = cv2.cvtColor(img1, cv2.COLOR_RGB2GRAY)
        gray2 = cv2.cvtColor(img2, cv2.COLOR_RGB2GRAY)
        
        # Calculate SSIM
        score, diff = ssim(gray1, gray2, full=True)
        
        # Convert to 0-100 scale
        similarity = score * 100
        
        # Create diff mask (areas with low similarity)
        diff_mask = ((1 - diff) * 255).astype(np.uint8)
        
        return similarity, diff_mask
    
    def create_diff_overlay(
        self,
        figma_img: np.ndarray,
        website_img: np.ndarray,
        diff_mask: np.ndarray,
        threshold: int = 30
    ) -> np.ndarray:
        """
        Create an overlay image highlighting differences.
        
        Args:
            figma_img: Figma screenshot as numpy array
            website_img: Website screenshot as numpy array
            diff_mask: Difference mask from SSIM
            threshold: Minimum difference to highlight (0-255)
            
        Returns:
            Overlay image with differences highlighted in red
        """
        # Create output image (copy of website)
        overlay = website_img.copy()
        
        # Find areas with significant differences
        significant_diff = diff_mask > threshold
        
        # Highlight differences in semi-transparent red
        red_overlay = overlay.copy()
        red_overlay[significant_diff] = [255, 0, 0]  # Red
        
        # Blend with original (50% opacity for red areas)
        alpha = 0.5
        overlay[significant_diff] = (
            alpha * red_overlay[significant_diff] + 
            (1 - alpha) * overlay[significant_diff]
        ).astype(np.uint8)
        
        return overlay
    
    def create_comparison_image(
        self,
        figma_path: str,
        website_path: str,
        output_path: str,
        figma_dims: Dict[str, int],
        website_dims: Dict[str, int],
        viewport: str
    ) -> Dict[str, Any]:
        """
        Create a comprehensive comparison image.
        
        Generates a side-by-side view:
        [Figma Design] | [Website] | [Diff Overlay]
        
        Returns:
            Dict with comparison results
        """
        print(f"\n  🔍 Comparing {viewport} screenshots...")
        
        # Normalize images
        figma_arr, website_arr, scale = self.normalize_images(
            figma_path, website_path, figma_dims, website_dims
        )
        
        # Calculate similarity
        similarity, diff_mask = self.calculate_similarity(figma_arr, website_arr)
        print(f"     📊 Similarity Score: {similarity:.1f}%")
        
        # Create diff overlay
        overlay = self.create_diff_overlay(figma_arr, website_arr, diff_mask)
        
        # Count different pixels
        significant_diff = diff_mask > 30
        diff_percentage = (np.sum(significant_diff) / significant_diff.size) * 100
        print(f"     📍 Pixels with differences: {diff_percentage:.1f}%")
        
        # Create side-by-side comparison
        h, w = figma_arr.shape[:2]
        padding = 20
        label_height = 40
        
        # Create canvas
        canvas_w = (w * 3) + (padding * 4)
        canvas_h = h + label_height + (padding * 2)
        canvas = np.ones((canvas_h, canvas_w, 3), dtype=np.uint8) * 240  # Light gray bg
        
        # Place images
        y_offset = label_height + padding
        
        # Figma (left)
        x1 = padding
        canvas[y_offset:y_offset+h, x1:x1+w] = figma_arr
        
        # Website (center)
        x2 = padding * 2 + w
        canvas[y_offset:y_offset+h, x2:x2+w] = website_arr
        
        # Diff overlay (right)
        x3 = padding * 3 + w * 2
        canvas[y_offset:y_offset+h, x3:x3+w] = overlay
        
        # Convert to PIL for text
        canvas_pil = Image.fromarray(canvas)
        draw = ImageDraw.Draw(canvas_pil)
        
        # Try to use a font, fall back to default
        try:
            font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf", 20)
            small_font = ImageFont.truetype("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", 14)
        except:
            font = ImageFont.load_default()
            small_font = font
        
        # Add labels
        draw.text((x1 + w//2 - 60, 10), "FIGMA DESIGN", fill=(0, 0, 0), font=font)
        draw.text((x2 + w//2 - 40, 10), "WEBSITE", fill=(0, 0, 0), font=font)
        draw.text((x3 + w//2 - 80, 10), "DIFFERENCES", fill=(255, 0, 0), font=font)
        
        # Add similarity score
        score_text = f"Similarity: {similarity:.1f}%"
        draw.text((canvas_w - 150, canvas_h - 30), score_text, fill=(0, 100, 0), font=small_font)
        
        # Save
        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
        canvas_pil.save(output_path)
        print(f"     ✓ Saved comparison: {output_path}")
        
        # Detect specific differences
        differences = self._detect_differences(figma_arr, website_arr, diff_mask, viewport)
        
        return {
            "viewport": viewport,
            "similarity_score": similarity,
            "diff_percentage": diff_percentage,
            "comparison_image": output_path,
            "differences": differences,
            "scale_applied": scale
        }
    
    def _detect_differences(
        self,
        figma_arr: np.ndarray,
        website_arr: np.ndarray,
        diff_mask: np.ndarray,
        viewport: str
    ) -> List[Dict[str, Any]]:
        """
        Detect and categorize specific differences.
        
        Returns:
            List of detected differences with details
        """
        differences = []
        
        # 1. Check overall color difference
        figma_mean = np.mean(figma_arr, axis=(0, 1))
        website_mean = np.mean(website_arr, axis=(0, 1))
        color_diff = np.linalg.norm(figma_mean - website_mean)
        
        if color_diff > 10:
            differences.append({
                "category": "colors",
                "severity": "Medium" if color_diff < 30 else "High",
                "title": "Color scheme differs",
                "description": f"Average color difference detected (delta: {color_diff:.1f})",
                "viewport": viewport
            })
        
        # 2. Check for significant regions of difference
        # Find contours in diff mask
        _, binary = cv2.threshold(diff_mask, 50, 255, cv2.THRESH_BINARY)
        contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        
        # Filter significant contours (larger than 1% of image)
        min_area = (figma_arr.shape[0] * figma_arr.shape[1]) * 0.01
        significant_regions = [c for c in contours if cv2.contourArea(c) > min_area]
        
        if len(significant_regions) > 0:
            differences.append({
                "category": "layout",
                "severity": "High" if len(significant_regions) > 5 else "Medium",
                "title": f"Layout differences in {len(significant_regions)} regions",
                "description": f"Found {len(significant_regions)} areas with significant visual differences",
                "viewport": viewport,
                "regions_count": len(significant_regions)
            })
        
        # 3. Check edges/borders
        figma_edges = cv2.Canny(cv2.cvtColor(figma_arr, cv2.COLOR_RGB2GRAY), 50, 150)
        website_edges = cv2.Canny(cv2.cvtColor(website_arr, cv2.COLOR_RGB2GRAY), 50, 150)
        edge_diff = np.abs(figma_edges.astype(float) - website_edges.astype(float))
        edge_diff_percentage = np.mean(edge_diff) / 255 * 100
        
        if edge_diff_percentage > 5:
            differences.append({
                "category": "structure",
                "severity": "Medium",
                "title": "Element borders/edges differ",
                "description": f"Edge structure differs by {edge_diff_percentage:.1f}%",
                "viewport": viewport
            })
        
        return differences
    
    def compare_all_viewports(
        self,
        figma_screenshots: Dict[str, str],
        website_screenshots: Dict[str, str],
        figma_dims: Dict[str, Dict[str, int]],
        website_dims: Dict[str, Dict[str, int]],
        execution_id: str
    ) -> Dict[str, Any]:
        """
        Compare all viewports and generate comprehensive results.
        
        Returns:
            Complete comparison results
        """
        results = {
            "comparisons": {},
            "all_differences": [],
            "viewport_scores": {},
            "overall_score": 0.0
        }
        
        viewports = set(figma_screenshots.keys()) & set(website_screenshots.keys())
        
        for viewport in viewports:
            output_path = f"{self.output_dir}/comparison_{viewport}_{execution_id}.png"
            
            comparison = self.create_comparison_image(
                figma_screenshots[viewport],
                website_screenshots[viewport],
                output_path,
                figma_dims.get(viewport, {}),
                website_dims.get(viewport, {}),
                viewport
            )
            
            results["comparisons"][viewport] = comparison
            results["all_differences"].extend(comparison["differences"])
            results["viewport_scores"][viewport] = comparison["similarity_score"]
        
        # Calculate overall score (average of viewports)
        if results["viewport_scores"]:
            results["overall_score"] = sum(results["viewport_scores"].values()) / len(results["viewport_scores"])
        
        return results