from abc import ABC, abstractmethod

from PIL import Image, ImageFilter


def resize_for_upload(image: Image.Image, max_dimension: int = 1536) -> Image.Image:
    """Downscale so the longest edge is at most max_dimension, keeping aspect ratio."""
    width, height = image.size
    longest_edge = max(width, height)
    if longest_edge <= max_dimension:
        return image
    scale = max_dimension / longest_edge
    new_size = (round(width * scale), round(height * scale))
    return image.resize(new_size, Image.LANCZOS)


class Segmenter(ABC):
    """Produces object masks from an image and point prompts."""

    @abstractmethod
    def encode(self, image: Image.Image) -> dict:
        """Validate/warm up the model for this image; return small JSON-safe metadata.

        Must not return tensors or other GPU-resident objects: under ZeroGPU,
        encode() and decode() may run in different ephemeral worker
        processes, and only the host (which has no real GPU) can persist
        data between them - so decode() re-derives the embedding from the
        original image instead of reusing anything encode() returns.
        """

    @abstractmethod
    def decode(self, image: Image.Image, points: list[tuple[int, int, bool]]) -> Image.Image:
        """Encode the image and decode it with the given point prompts; return an 'L' mode mask."""


import numpy as np


class Sam2ModelSegmenter(Segmenter):
    """Wraps SAM2 hiera-base-plus via the `sam2` package for server-side inference."""

    def __init__(self, hf_model_id: str = "facebook/sam2.1-hiera-base-plus", device: str = "cuda"):
        # Do NOT build the model here: __init__ runs at Space startup, which
        # is outside any @spaces.GPU-decorated call. ZeroGPU's CUDA
        # emulation mode active at that point cannot correctly satisfy
        # SAM2's Hydra-based instantiation of MemoryAttention (it raises
        # "Expected all tensors to be on the same device ... cuda:0
        # (ZeroGPU) and cpu" regardless of whether "cpu" or "cuda" is
        # requested). Build lazily on first encode(), which - once this
        # method is wrapped with @spaces.GPU - runs inside a real,
        # physically-attached GPU worker instead of the emulation.
        self._hf_model_id = hf_model_id
        self._device = device
        self._predictor = None

    def _ensure_loaded(self) -> None:
        if self._predictor is not None:
            return
        from sam2.build_sam import build_sam2_hf
        from sam2.sam2_image_predictor import SAM2ImagePredictor

        sam2_model = build_sam2_hf(self._hf_model_id, device=self._device)
        self._predictor = SAM2ImagePredictor(sam2_model)

    def encode(self, image: Image.Image) -> dict:
        self._ensure_loaded()
        return {"image_size": image.size}

    def decode(
        self, image: Image.Image, points: list[tuple[int, int, bool]], edge_padding: int = 6
    ) -> Image.Image:
        self._ensure_loaded()
        image_array = np.array(image.convert("RGB"))
        self._predictor.set_image(image_array)
        point_coords = np.array([[x, y] for x, y, _ in points])
        point_labels = np.array([1 if positive else 0 for _, _, positive in points])
        masks, scores, _ = self._predictor.predict(
            point_coords=point_coords,
            point_labels=point_labels,
            multimask_output=True,
        )
        best_mask = masks[int(np.argmax(scores))]
        mask_image = Image.fromarray((best_mask * 255).astype("uint8"), mode="L")
        mask_image = mask_image.resize(image.size)
        # SAM2's boundary tends to sit a few pixels inside the object's true
        # edge rather than exactly on it, leaving a visible sliver of the
        # object outside the selection - padding it out here means the mask
        # the user sees (and that gets sent for removal) already comfortably
        # covers the object, instead of relying on them to notice and patch
        # small gaps with extra taps/brush strokes.
        if edge_padding > 0:
            mask_image = mask_image.filter(ImageFilter.MaxFilter(edge_padding * 2 + 1))
        return mask_image