from __future__ import annotations

import numpy as np
from PIL import Image, ImageDraw, ImageFont

from src.config.endpoints import BBOX_NORMALIZATION_SCALE
from src.schemas.detection import LocalizedFinding


def draw_view_localizations(
    image: Image.Image,
    image_index: int,
    findings: list[LocalizedFinding],
) -> Image.Image:
    annotated = image.copy()
    draw = ImageDraw.Draw(annotated)
    font = ImageFont.load_default()

    width, height = annotated.size

    for finding in findings:
        for view_localization in finding.localizations:
            if view_localization.image_index != image_index:
                continue

            for localization in view_localization.boxes:
                y0, x0, y1, x1 = localization.box_2d

                left = int(x0 / BBOX_NORMALIZATION_SCALE * width)
                top = int(y0 / BBOX_NORMALIZATION_SCALE * height)
                right = int(x1 / BBOX_NORMALIZATION_SCALE * width)
                bottom = int(y1 / BBOX_NORMALIZATION_SCALE * height)

                draw.rectangle(
                    [(left, top), (right, bottom)],
                    outline="red",
                    width=max(2, width // 256),
                )

                label = localization.label
                text_bbox = draw.textbbox((left, top), label, font=font)
                text_width = text_bbox[2] - text_bbox[0]
                text_height = text_bbox[3] - text_bbox[1]
                label_top = max(0, top - text_height - 6)

                draw.rectangle(
                    [
                        (left, label_top),
                        (left + text_width + 6, label_top + text_height + 6),
                    ],
                    fill="black",
                )
                draw.text(
                    (left + 3, label_top + 3),
                    label,
                    fill="white",
                    font=font,
                )

    return annotated


def draw_segmentation_overlay(
    base_images: list[Image.Image],
    masks: list,
) -> list[Image.Image]:
    """Overlay polygon segmentation masks onto images, matching the MedSAM notebook style.

    Draws a semi-transparent yellow polygon fill + outline for each successful
    MaskResult.  Falls back to a pixel-level mask blend when no polygon is available.
    Also draws the bounding box that prompted the segmentation.
    """
    result = [img.convert("RGBA") for img in base_images]

    for mask_result in masks:
        if mask_result.status != "success":
            continue
        idx = mask_result.image_index
        if idx >= len(result):
            continue

        img = result[idx]
        W, H = img.size
        overlay = Image.new("RGBA", (W, H), (0, 0, 0, 0))
        draw = ImageDraw.Draw(overlay)

        if mask_result.polygon and len(mask_result.polygon) >= 3:
            pts = [(x, y) for x, y in mask_result.polygon]
            # Yellow fill (251, 252, 30) @ ~60% opacity, bright outline
            draw.polygon(pts, fill=(251, 252, 30, 153), outline=(251, 220, 0, 230))
        elif mask_result.mask is not None:
            # Pixel-level fallback: blend yellow where mask is 1
            base_np = np.array(img, dtype=np.float32)
            mask = mask_result.mask
            if mask.shape != (H, W):
                mask_pil = Image.fromarray((mask * 255).astype(np.uint8)).resize(
                    (W, H), Image.NEAREST
                )
                mask = (np.array(mask_pil) > 127).astype(np.uint8)
            yellow = np.array([251, 252, 30, 255], dtype=np.float32)
            alpha = 0.6
            mask_bool = mask.astype(bool)
            base_np[mask_bool, :3] = base_np[mask_bool, :3] * (1 - alpha) + yellow[:3] * alpha
            img = Image.fromarray(base_np.clip(0, 255).astype(np.uint8), mode="RGBA")
            overlay = Image.new("RGBA", (W, H), (0, 0, 0, 0))
            draw = ImageDraw.Draw(overlay)

        # Bounding box in blue (matching the notebook)
        y0_n, x0_n, y1_n, x1_n = mask_result.bbox_normalized
        x0 = int(x0_n / BBOX_NORMALIZATION_SCALE * W)
        y0 = int(y0_n / BBOX_NORMALIZATION_SCALE * H)
        x1 = int(x1_n / BBOX_NORMALIZATION_SCALE * W)
        y1 = int(y1_n / BBOX_NORMALIZATION_SCALE * H)
        lw = max(2, W // 256)
        draw.rectangle([(x0, y0), (x1, y1)], outline=(0, 100, 255, 255), width=lw)

        result[idx] = Image.alpha_composite(img, overlay)

    return [img.convert("RGB") for img in result]