brain.py

import os
import json
import typing_extensions
from dotenv import load_dotenv
import google.generativeai as genai
from PIL import Image, ImageDraw, ImageFont

# Load environment variables
load_dotenv()

# Securely configure the API
api_key = os.getenv("GOOGLE_API_KEY")
if api_key:
    genai.configure(api_key=api_key)
    # Flash model: Low latency, high reasoning for visual tasks
    model = genai.GenerativeModel('gemini-2.0-flash-001')
else:
    model = None
    print("⚠️ Warning: GOOGLE_API_KEY not found in environment.")

# Strict Type Definition for the AI Response
class Detection(typing_extensions.TypedDict):
    label: str
    explanation: str
    severity: str 
    box_2d: list[int] 

def analyze_screenshot(image_path: str):
    """
    Sends visual data to Gemini to identify deceptive patterns.
    """
    print("🧠 Brain: Analyzing UI semantics and visual hierarchy...")
    
    if not model:
        return []

    try:
        img = Image.open(image_path)
    except Exception as e:
        print(f"❌ Brain: Could not load image. {e}")
        return []

    prompt = """
    You are a world-class UI/UX Auditor specializing in Ethical Design.
    Your job is to detect 'Dark Patterns' in the provided website screenshot.
    
    Target Categories:
    1. Urgency/Scarcity (e.g., "Only 2 rooms left", countdowns)
    2. Misdirection (e.g., Confusing button styles, hidden costs)
    3. Social Proof (e.g., "150 people viewed this")
    4. Obstruction (e.g., Hard to cancel)

    Output STRICT JSON format:
    [
      {
        "label": "Short Title",
        "explanation": "Why is this deceptive?",
        "severity": "High" or "Medium",
        "box_2d": [ymin, xmin, ymax, xmax] (0-1000 scale)
      }
    ]
    
    If clean, return [].
    """

    try:
        response = model.generate_content(
            [prompt, img],
            generation_config={"response_mime_type": "application/json"}
        )
        detections = json.loads(response.text)
        print(f"🧠 Brain: Identified {len(detections)} issues.")
        return detections
        
    except Exception as e:
        print(f"❌ Brain: Cognitive failure (API Error): {e}")
        return []

def draw_boxes(image_path: str, detections: list, output_path: str = "annotated_evidence.png"):
    """
    Draws professional bounding boxes with transparency.
    """
    if not detections:
        return image_path

    print("🎨 Artist: Rendering detection overlays...")
    
    image = Image.open(image_path).convert("RGBA")
    # Create a separate layer for transparency
    overlay = Image.new("RGBA", image.size, (255, 255, 255, 0))
    draw = ImageDraw.Draw(overlay)
    
    width, height = image.size

    for item in detections:
        ymin, xmin, ymax, xmax = item['box_2d']
        
        # Convert normalized coordinates to pixels
        box = [
            xmin / 1000 * width,
            ymin / 1000 * height,
            xmax / 1000 * width,
            ymax / 1000 * height
        ]
        
        # Color coding: Red for High, Orange for Medium
        color = (255, 0, 0, 80) if item.get("severity") == "High" else (255, 165, 0, 80)
        outline = (255, 0, 0, 255) if item.get("severity") == "High" else (255, 165, 0, 255)

        # Draw filled box
        draw.rectangle(box, fill=color, outline=outline, width=4)
        
        # Label Badge
        text_x, text_y = box[0], max(0, box[1] - 30)
        label_text = f" {item['label']} "
        
        try:
            font = ImageFont.truetype("arial.ttf", 22)
        except:
            font = ImageFont.load_default()

        # Text background
        draw.rectangle([text_x, text_y, text_x + 200, text_y + 30], fill=outline)
        draw.text((text_x + 5, text_y + 2), label_text, fill="white", font=font)

    # Composite and save
    final_image = Image.alpha_composite(image, overlay).convert("RGB")
    final_image.save(output_path)
    print(f"✅ Artist: Evidence compiled at {output_path}")
    return output_path