"""Lazy-loaded FastSAM wrapper for click-to-annotate.

The user clicks a pixel on the active page; we run FastSAM's
"segment-everything" pass over the image, then pick the smallest mask whose
binary pixel value at the click is above threshold (smallest = most local
object, the one the user most likely meant).

The model file (~74 MB for FastSAM-s) auto-downloads on first use.
"""
from __future__ import annotations

import base64
import io
import threading
from typing import Optional

import numpy as np
from PIL import Image

_MODEL_LOCK = threading.Lock()
_MODEL = None
_MODEL_NAME = "FastSAM-s.pt"  # ~74 MB, fast enough for interactive use


def _get_model():
    global _MODEL
    if _MODEL is not None:
        return _MODEL
    with _MODEL_LOCK:
        if _MODEL is None:
            from ultralytics import FastSAM  # imported lazily to keep startup quick
            _MODEL = FastSAM(_MODEL_NAME)
    return _MODEL


def segment_at_point(
    image_b64: str, x: int, y: int, *,
    imgsz: int = 1024, conf: float = 0.4, iou: float = 0.9,
) -> Optional[dict]:
    """Return ``{"bbox_px": [x1, y1, x2, y2]}`` for the FastSAM mask under
    the click, or ``None`` if no mask covers that pixel.

    Coordinates are in the natural image frame (the same frame the
    frontend uses when sending the click).
    """
    raw = base64.b64decode(image_b64)
    img = Image.open(io.BytesIO(raw)).convert("RGB")
    arr = np.array(img)
    img_h, img_w = arr.shape[:2]

    if not (0 <= x < img_w and 0 <= y < img_h):
        return None

    model = _get_model()
    results = model.predict(
        arr,
        device="cpu",
        retina_masks=True,
        imgsz=imgsz,
        conf=conf,
        iou=iou,
        verbose=False,
    )
    if not results:
        return None
    result = results[0]
    if result.masks is None or len(result.masks.data) == 0:
        return None

    masks = result.masks.data.cpu().numpy() if hasattr(result.masks.data, "cpu") else np.asarray(result.masks.data)
    # Masks are (N, H, W) floats in [0, 1] — may not be at the original resolution if
    # retina_masks is unavailable for this build; rescale the click accordingly.
    mh, mw = masks.shape[1], masks.shape[2]
    mx = int(round(x * mw / img_w))
    my = int(round(y * mh / img_h))
    mx = max(0, min(mw - 1, mx))
    my = max(0, min(mh - 1, my))

    candidates: list[tuple[int, list[int]]] = []
    for i in range(masks.shape[0]):
        if masks[i, my, mx] <= 0.5:
            continue
        ys, xs = np.where(masks[i] > 0.5)
        if xs.size == 0:
            continue
        # bbox in mask-coordinates, then scaled to image-coordinates
        x1 = int(round(xs.min() * img_w / mw))
        y1 = int(round(ys.min() * img_h / mh))
        x2 = int(round((xs.max() + 1) * img_w / mw))
        y2 = int(round((ys.max() + 1) * img_h / mh))
        x1 = max(0, min(img_w, x1)); y1 = max(0, min(img_h, y1))
        x2 = max(0, min(img_w, x2)); y2 = max(0, min(img_h, y2))
        if x2 <= x1 or y2 <= y1:
            continue
        area = int(xs.size)
        candidates.append((area, [x1, y1, x2, y2]))

    if not candidates:
        return None

    # The smallest containing mask is almost always the most local / most
    # precise object the user pointed at.
    candidates.sort(key=lambda t: t[0])
    return {"bbox_px": candidates[0][1]}