# Copyright (c) 2025 CoReason, Inc.
#
# This software is proprietary and dual-licensed.
# Licensed under the Prosperity Public License 3.0 (the "License").
# A copy of the license is available at https://prosperitylicense.com/versions/3.0.0
# For details, see the LICENSE file.
# Commercial use beyond a 30-day trial requires a separate license.
#
# Source Code: https://github.com/CoReason-AI/coreason_navigator

import base64
import io
from typing import List

from PIL import Image, ImageDraw, ImageFont

from coreason_navigator.types import ElementBoundingBox
from coreason_navigator.utils.logger import logger


def draw_som_overlay(screenshot_base64: str, elements: List[ElementBoundingBox]) -> str:
    """
    Draws a Set-of-Marks (SoM) overlay on the screenshot.

    Args:
        screenshot_base64: The base64 encoded original screenshot.
        elements: List of interactive elements with bounding boxes.

    Returns:
        str: Base64 encoded string of the screenshot with the overlay.
    """
    try:
        # Decode base64 image
        image_data = base64.b64decode(screenshot_base64)
        image = Image.open(io.BytesIO(image_data))

        # Ensure image is in RGB mode
        if image.mode != "RGB":
            image = image.convert("RGB")

        draw = ImageDraw.Draw(image)

        # Load font (try default, or use a specific one if available)
        # Attempt to use a larger default font
        font = ImageFont.load_default()
        # Note: load_default() might be too small.
        # Ideally we'd load a truetype font but we can't guarantee system fonts.
        # For now, let's stick with default or try to scale if possible?
        # PIL default font is bitmap and fixed size.
        # We will use it for now as a safe fallback.

        for index, element in enumerate(elements):
            # Draw bounding box
            # Use red color for visibility
            box = [element.x, element.y, element.x + element.width, element.y + element.height]
            draw.rectangle(box, outline="red", width=2)

            # Draw label (numeric tag)
            # We use the index as the tag for LLM reference
            label = str(index)

            # Position label: try top-left of the box, but ensure it's within image bounds
            # Draw a filled rectangle background for the text to be readable
            text_bbox = draw.textbbox((0, 0), label, font=font)
            text_width = text_bbox[2] - text_bbox[0]
            text_height = text_bbox[3] - text_bbox[1]

            # Label position
            label_x = element.x
            label_y = element.y - text_height - 4

            if label_y < 0:
                label_y = element.y + 4  # Move inside/below if off-screen top

            # Background for label
            draw.rectangle(
                [label_x, label_y, label_x + text_width + 4, label_y + text_height + 4], fill="red", outline="red"
            )

            # Text
            draw.text((label_x + 2, label_y + 2), label, fill="white", font=font)

        # Encode back to base64
        buffered = io.BytesIO()
        image.save(buffered, format="JPEG", quality=75)
        return base64.b64encode(buffered.getvalue()).decode("utf-8")

    except Exception:
        logger.exception("Failed to generate SoM overlay")
        # Return original image if overlay fails to be safe
        return screenshot_base64
