agents/agent_3_difference_analyzer.py

"""
Agent 3: Difference Analyzer (Hybrid Approach)
Uses HF Vision Model + Screenshot Analysis to detect visual differences
Detects layout, color, spacing, typography, and other visual differences
"""

from typing import Dict, Any, List
from state_schema import WorkflowState, VisualDifference
import os
from PIL import Image
import numpy as np


class ScreenshotComparator:
    """Compares Figma and website screenshots for visual differences."""
    
    def __init__(self):
        """Initialize the comparator."""
        self.differences = []
    
    def compare_screenshots(self, figma_path: str, website_path: str, viewport: str) -> List[VisualDifference]:
        """
        Compare Figma and website screenshots.
        
        Args:
            figma_path: Path to Figma screenshot
            website_path: Path to website screenshot
            viewport: Viewport name (desktop/mobile)
        
        Returns:
            List of detected visual differences
        """
        differences = []
        
        if not os.path.exists(figma_path) or not os.path.exists(website_path):
            return differences
        
        try:
            # Load images
            figma_img = Image.open(figma_path)
            website_img = Image.open(website_path)
            
            # Analyze different aspects
            differences.extend(self._analyze_layout(figma_img, website_img, viewport))
            differences.extend(self._analyze_colors(figma_img, website_img, viewport))
            differences.extend(self._analyze_structure(figma_img, website_img, viewport))
            
        except Exception as e:
            print(f"Error comparing screenshots: {str(e)}")
        
        return differences
    
    def _analyze_layout(self, figma_img: Image.Image, website_img: Image.Image, viewport: str) -> List[VisualDifference]:
        """Analyze layout differences."""
        differences = []
        
        figma_size = figma_img.size
        website_size = website_img.size
        
        # Check width differences
        if figma_size[0] != website_size[0]:
            width_diff_percent = abs(figma_size[0] - website_size[0]) / figma_size[0] * 100
            severity = "High" if width_diff_percent > 10 else "Medium"
            
            diff = VisualDifference(
                category="layout",
                severity=severity,
                issue_id="1.1",
                title="Container width differs",
                description=f"Design: {figma_size[0]}px vs Website: {website_size[0]}px ({width_diff_percent:.1f}% difference)",
                design_value=str(figma_size[0]),
                website_value=str(website_size[0]),
                viewport=viewport,
                confidence=0.95,
                detection_method="screenshot_comparison"
            )
            differences.append(diff)
        
        # Check height differences
        if figma_size[1] != website_size[1]:
            height_diff_percent = abs(figma_size[1] - website_size[1]) / figma_size[1] * 100
            severity = "High" if height_diff_percent > 10 else "Medium"
            
            diff = VisualDifference(
                category="layout",
                severity=severity,
                issue_id="1.2",
                title="Page height differs",
                description=f"Design: {figma_size[1]}px vs Website: {website_size[1]}px ({height_diff_percent:.1f}% difference)",
                design_value=str(figma_size[1]),
                website_value=str(website_size[1]),
                viewport=viewport,
                confidence=0.95,
                detection_method="screenshot_comparison"
            )
            differences.append(diff)
        
        return differences
    
    def _analyze_colors(self, figma_img: Image.Image, website_img: Image.Image, viewport: str) -> List[VisualDifference]:
        """Analyze color differences."""
        differences = []
        
        try:
            # Convert to RGB
            figma_rgb = figma_img.convert('RGB')
            website_rgb = website_img.convert('RGB')
            
            # Resize to same size for comparison (use smaller size for performance)
            compare_size = (400, 300)
            figma_resized = figma_rgb.resize(compare_size)
            website_resized = website_rgb.resize(compare_size)
            
            # Convert to numpy arrays
            figma_array = np.array(figma_resized, dtype=np.float32)
            website_array = np.array(website_resized, dtype=np.float32)
            
            # Calculate color difference (mean absolute difference)
            color_diff = np.mean(np.abs(figma_array - website_array))
            
            # If significant color difference, flag it
            if color_diff > 15:
                severity = "High" if color_diff > 40 else "Medium"
                
                diff = VisualDifference(
                    category="colors",
                    severity=severity,
                    issue_id="3.1",
                    title="Color scheme differs significantly",
                    description=f"Significant color difference detected (delta: {color_diff:.1f})",
                    design_value="Design colors",
                    website_value="Website colors",
                    viewport=viewport,
                    confidence=0.8,
                    detection_method="pixel_analysis"
                )
                differences.append(diff)
        
        except Exception as e:
            pass
        
        return differences
    
    def _analyze_structure(self, figma_img: Image.Image, website_img: Image.Image, viewport: str) -> List[VisualDifference]:
        """Analyze structural/layout differences."""
        differences = []
        
        try:
            # Convert to grayscale for edge detection
            figma_gray = figma_img.convert('L')
            website_gray = website_img.convert('L')
            
            # Resize to same size
            compare_size = (400, 300)
            figma_resized = figma_gray.resize(compare_size)
            website_resized = website_gray.resize(compare_size)
            
            # Convert to numpy arrays
            figma_array = np.array(figma_resized, dtype=np.float32)
            website_array = np.array(website_resized, dtype=np.float32)
            
            # Calculate structural difference (MSE)
            mse = np.mean((figma_array - website_array) ** 2)
            
            # Normalize MSE to 0-100 scale
            structural_diff = min(100, mse / 255)
            
            if structural_diff > 10:
                severity = "High" if structural_diff > 30 else "Medium"
                
                diff = VisualDifference(
                    category="layout",
                    severity=severity,
                    issue_id="1.3",
                    title="Layout structure differs",
                    description=f"Visual structure difference detected (score: {structural_diff:.1f})",
                    design_value="Design layout",
                    website_value="Website layout",
                    viewport=viewport,
                    confidence=0.75,
                    detection_method="structural_analysis"
                )
                differences.append(diff)
        
        except Exception as e:
            pass
        
        return differences


class DifferenceAnalyzer:
    """
    Agent 3: Difference Analyzer
    Analyzes visual differences between Figma designs and website implementations
    """
    
    def __init__(self):
        """Initialize the analyzer."""
        self.comparator = ScreenshotComparator()
    
    def analyze_differences(self, state: WorkflowState) -> WorkflowState:
        """
        Analyze visual differences between Figma and website screenshots.
        
        Args:
            state: Current workflow state
        
        Returns:
            Updated state with analysis results
        """
        print("\n🔍 Agent 3: Difference Analyzer - Analyzing Visual Differences...")
        
        try:
            all_differences = []
            
            # Compare screenshots for each viewport
            for viewport in ["desktop", "mobile"]:
                figma_key = f"{viewport}"
                website_key = f"{viewport}"
                
                figma_path = state.get("figma_screenshots", {}).get(figma_key)
                website_path = state.get("website_screenshots", {}).get(website_key)
                
                if figma_path and website_path:
                    print(f"  📊 Comparing {viewport} screenshots...")
                    
                    differences = self.comparator.compare_screenshots(
                        figma_path,
                        website_path,
                        viewport
                    )
                    
                    all_differences.extend(differences)
                    print(f"    ✓ Found {len(differences)} differences")
                else:
                    print(f"  ⚠️  Missing screenshots for {viewport}")
            
            # Calculate similarity score
            total_differences = len(all_differences)
            high_severity = len([d for d in all_differences if d.severity == "High"])
            medium_severity = len([d for d in all_differences if d.severity == "Medium"])
            low_severity = len([d for d in all_differences if d.severity == "Low"])
            
            # Similarity score: 100 - (differences weighted by severity)
            severity_weight = (high_severity * 10) + (medium_severity * 5) + (low_severity * 1)
            similarity_score = max(0, 100 - severity_weight)
            
            state["visual_differences"] = [d.to_dict() if hasattr(d, "to_dict") else d for d in all_differences]
            state["similarity_score"] = similarity_score
            state["status"] = "analysis_complete"
            
            print(f"\n  📈 Analysis Summary:")
            print(f"    - Total differences: {total_differences}")
            print(f"    - High severity: {high_severity}")
            print(f"    - Medium severity: {medium_severity}")
            print(f"    - Low severity: {low_severity}")
            print(f"    - Similarity score: {similarity_score:.1f}/100")
            
            return state
            
        except Exception as e:
            print(f"  ❌ Error analyzing differences: {str(e)}")
            import traceback
            traceback.print_exc()
            state["status"] = "analysis_failed"
            state["error_message"] = f"Agent 3 Error: {str(e)}"
            return state


def agent_3_node(state: Dict) -> Dict:
    """
    LangGraph node for Agent 3 (Difference Analyzer).
    
    Args:
        state: Current workflow state
    
    Returns:
        Updated state
    """
    # Convert dict to WorkflowState if needed
    if isinstance(state, dict):
        workflow_state = WorkflowState(**state)
    else:
        workflow_state = state
    
    # Create analyzer and analyze differences
    analyzer = DifferenceAnalyzer()
    updated_state = analyzer.analyze_differences(workflow_state)
    
    # Convert back to dict for LangGraph
    return updated_state.__dict__