modules/element_annotation.py

import os
import logging
from typing import List, Dict, Tuple
from PIL import Image, ImageDraw, ImageFont
from modules.config import ANNOTATION_DIR, FONT_SIZE, JPEG_QUALITY, ANNOTATE_TEXT, DRAW_TEXT_BOXES, DRAW_ARROWS, ARROW_WIDTH, LABEL_BORDER_WIDTH

logger = logging.getLogger(__name__)

def get_text_color(image: Image.Image, x: int, y: int, size: int) -> tuple:
    """Determine the best text color (black or white) based on background brightness."""
    # Get a small region around the point to analyze brightness
    region = image.crop((
        max(0, x - size//2), 
        max(0, y - size//2), 
        min(image.width, x + size//2), 
        min(image.height, y + size//2)
    ))
    
    # Convert to grayscale and calculate average brightness
    gray_region = region.convert("L")
    brightness = sum(gray_region.getdata()) / len(gray_region.getdata())
    
    # Return white for dark backgrounds, black for light backgrounds
    return (255, 255, 255) if brightness < 128 else (0, 0, 0)

def check_label_collision(rect1, all_rects):
    """Check if a rectangle collides with any other rectangles."""
    x1, y1, x2, y2 = rect1
    for r in all_rects:
        rx1, ry1, rx2, ry2 = r
        # Check if rectangles overlap
        if not (x2 < rx1 or x1 > rx2 or y2 < ry1 or y1 > ry2):
            return True
    return False

def find_closest_point_on_bbox(bbox: Tuple[int, int, int, int], point: Tuple[int, int]) -> Tuple[int, int]:
    """Find the closest point on the bounding box to the given point."""
    x_min, y_min, x_max, y_max = bbox
    px, py = point
    
    # Check if point is inside the bbox
    if x_min <= px <= x_max and y_min <= py <= y_max:
        # Find the closest edge
        distances = [
            (abs(px - x_min), (x_min, py)),  # left edge
            (abs(px - x_max), (x_max, py)),  # right edge
            (abs(py - y_min), (px, y_min)),  # top edge
            (abs(py - y_max), (px, y_max))   # bottom edge
        ]
        return min(distances, key=lambda d: d[0])[1]
    
    # If point is outside the bbox, find the closest point
    x = max(x_min, min(px, x_max))
    y = max(y_min, min(py, y_max))
    return (x, y)

def find_closest_point_on_label(label_box: Tuple[int, int, int, int], point: Tuple[int, int]) -> Tuple[int, int]:
    """Find the closest point on the label box to the given point."""
    # Same logic as find_closest_point_on_bbox, but for label boxes
    return find_closest_point_on_bbox(label_box, point)

def find_non_colliding_position(text_x, text_y, text_width, text_height, bbox, elements, image_width, image_height) -> Tuple[int, int]:
    """Find a position for the code label that doesn't collide with other elements."""
    label_padding = 5  # Increased padding
    padding = label_padding + LABEL_BORDER_WIDTH
    
    # Original position (above the bbox)
    positions = [
        # Top (default)
        (text_x, max(0, text_y)),
        # Bottom
        (text_x, min(image_height - text_height - padding * 2, bbox[3] + padding)),
        # Left
        (max(0, bbox[0] - text_width - padding * 2), (bbox[1] + bbox[3]) // 2 - text_height // 2),
        # Right
        (min(image_width - text_width - padding * 2, bbox[2] + padding), (bbox[1] + bbox[3]) // 2 - text_height // 2),
        # Top-left corner
        (max(0, bbox[0] - text_width - padding), max(0, bbox[1] - text_height - padding)),
        # Top-right corner
        (min(image_width - text_width - padding, bbox[2] + padding), max(0, bbox[1] - text_height - padding)),
        # Bottom-left corner
        (max(0, bbox[0] - text_width - padding), min(image_height - text_height - padding, bbox[3] + padding)),
        # Bottom-right corner
        (min(image_width - text_width - padding, bbox[2] + padding), min(image_height - text_height - padding, bbox[3] + padding))
    ]
    
    # Try standard positions first
    for pos_x, pos_y in positions:
        rect = (pos_x - label_padding, pos_y - label_padding, 
                pos_x + text_width + label_padding, pos_y + text_height + label_padding)
        if not check_label_collision(rect, elements):
            return pos_x, pos_y
    
    # If all standard positions collide, find the first non-colliding position
    # by increasing distance from the bbox
    for distance in range(10, 100, 10):
        for angle in range(0, 360, 45):
            import math
            rad = math.radians(angle)
            center_x = (bbox[0] + bbox[2]) // 2
            center_y = (bbox[1] + bbox[3]) // 2
            
            pos_x = int(center_x + distance * math.cos(rad))
            pos_y = int(center_y + distance * math.sin(rad))
            
            # Ensure within image bounds
            pos_x = max(0, min(image_width - text_width - padding * 2, pos_x))
            pos_y = max(0, min(image_height - text_height - padding * 2, pos_y))
            
            rect = (pos_x - label_padding, pos_y - label_padding, 
                    pos_x + text_width + label_padding, pos_y + text_height + label_padding)
            if not check_label_collision(rect, elements):
                return pos_x, pos_y
    
    # If all positions collide, return the original position
    return text_x, text_y

async def annotate_image(image: Image.Image, elements: List[Dict], filename: str) -> str:
    """Annotate the image with element codes."""
    try:
        # Create a copy of the image for annotation
        annotated = image.copy()
        draw = ImageDraw.Draw(annotated)
        image_width, image_height = image.size
        
        # Try to load a font, or use default
        try:
            font = ImageFont.truetype("Arial.ttf", FONT_SIZE)
        except IOError:
            font = ImageFont.load_default()
        
        # First, collect all text element bboxes to avoid placing labels over text
        text_element_bboxes = []
        for element in elements:
            if element.get("type") == "text" and "bbox_pixels" in element and len(element["bbox_pixels"]) == 4:
                x_min, y_min, x_max, y_max = element["bbox_pixels"]
                # Add some padding around text elements
                text_element_bboxes.append((x_min - 4, y_min - 4, x_max + 4, y_max + 4))
        
        # Keep track of label positions to avoid overlaps
        label_positions = text_element_bboxes.copy()
        
        # Draw each element code
        for element in elements:
            # Skip elements without required fields
            if "code" not in element or "bbox_pixels" not in element:
                continue
            
            is_text_element = element.get("type") == "text"
            
            # Skip text elements if annotation for text is disabled
            if is_text_element and not ANNOTATE_TEXT:
                # Still collect bounding box for drawing if enabled
                if DRAW_TEXT_BOXES and len(element["bbox_pixels"]) == 4:
                    x_min, y_min, x_max, y_max = element["bbox_pixels"]
                    text_color = get_text_color(image, x_min, y_min, FONT_SIZE * 2)
                    draw.rectangle([x_min, y_min, x_max, y_max], outline=text_color, width=2)
                continue
                
            code = element["code"]
            bbox = element["bbox_pixels"]
            
            if len(bbox) != 4:
                continue
                
            x_min, y_min, x_max, y_max = bbox
            
            # Calculate dimensions for the code label
            text_width = FONT_SIZE * len(code) * 0.6
            text_height = FONT_SIZE
            
            # Default position (top-left of the bounding box)
            text_x = x_min
            text_y = max(0, y_min - text_height - 4)
            
            # Find a non-colliding position for the label
            text_x, text_y = find_non_colliding_position(
                text_x, text_y, text_width, text_height, 
                bbox, label_positions, image_width, image_height
            )
            
            # Label position already added above
            
            # Get appropriate text color
            text_color = get_text_color(image, text_x, text_y, FONT_SIZE * 2)
            outline_color = (0, 0, 0) if text_color[0] > 128 else (255, 255, 255)
            
            # Create label box with padding
            label_padding = 5  # Increased padding
            label_box = (
                text_x - label_padding, 
                text_y - label_padding, 
                text_x + text_width + label_padding, 
                text_y + text_height + label_padding
            )
            
            # Store label box for collision avoidance (moved up before usage in arrow drawing)
            label_positions.append(label_box)
            
            # Draw text background
            draw.rectangle(
                label_box,
                fill=(*outline_color, 180)
            )
            
            # Draw border around the label
            draw.rectangle(
                label_box,
                outline=text_color,
                width=LABEL_BORDER_WIDTH
            )
            
            # Draw text
            draw.text((text_x, text_y), code, font=font, fill=text_color)
            
            # Draw bounding box
            draw.rectangle([x_min, y_min, x_max, y_max], outline=text_color, width=2)
            
            # Draw arrow from label to bounding box if enabled
            if DRAW_ARROWS:
                # Calculate center points of element bbox
                element_center_x = (x_min + x_max) / 2
                element_center_y = (y_min + y_max) / 2
                
                # Find closest point on the label box (where arrow starts)
                arrow_start = find_closest_point_on_label(
                    label_box, (int(element_center_x), int(element_center_y))
                )
                
                # Find closest point on the element bbox (where arrow ends)
                arrow_end = find_closest_point_on_bbox(
                    bbox, arrow_start
                )
                
                # Draw the arrow line (from label box to element)
                draw.line(
                    [arrow_start, arrow_end], 
                    fill=text_color, 
                    width=ARROW_WIDTH
                )
                
                # Calculate direction vector (now for arrow going TO element, not FROM label)
                dx = arrow_end[0] - arrow_start[0]
                dy = arrow_end[1] - arrow_start[1]
                
                # Normalize the direction vector
                length = (dx**2 + dy**2)**0.5
                if length > 0:
                    dx, dy = dx/length, dy/length
                    
                    # Calculate arrowhead points
                    arrow_size = max(ARROW_WIDTH * 3, 8)  # Scale with arrow width
                    
                    # Calculate perpendicular vectors for the arrowhead
                    perpx, perpy = -dy, dx
                    
                    # Calculate the three points of the arrowhead
                    point1 = arrow_end
                    point2 = (
                        int(arrow_end[0] - dx * arrow_size - perpx * arrow_size/2),
                        int(arrow_end[1] - dy * arrow_size - perpy * arrow_size/2)
                    )
                    point3 = (
                        int(arrow_end[0] - dx * arrow_size + perpx * arrow_size/2),
                        int(arrow_end[1] - dy * arrow_size + perpy * arrow_size/2)
                    )
                    
                    # Draw the filled arrowhead triangle
                    draw.polygon([point1, point2, point3], fill=text_color)
        
        # Save the annotated image
        output_path = os.path.join(ANNOTATION_DIR, filename)
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        annotated.save(output_path, "JPEG", quality=JPEG_QUALITY)
        logger.info(f"Annotated image saved to {output_path}")
        
        return output_path
        
    except Exception as e:
        logger.error(f"Error annotating image: {str(e)}")
        return ""