"""Floor/region segmentation (SAM) with a geometric fallback.

When a SAM checkpoint is configured (SAM_CHECKPOINT) and the ML stack is present,
this isolates the floor with point-prompted Segment Anything. Otherwise it
returns a robust geometric floor mask (region below the horizon, narrowing toward
the back wall). Either way the Shapely stage consumes a binary floor mask.
"""
from __future__ import annotations

from pathlib import Path

import cv2
import numpy as np

from ..config import settings
from .ml_runtime import ml_available


def floor_mask(image_path: Path, horizon_frac: float | None = None) -> np.ndarray:
    """Return a HxW uint8 mask (255 = floor) for the visible floor region."""
    img = cv2.imread(str(image_path))
    if img is None:
        raise ValueError(f"Could not read image for segmentation: {image_path}")
    h, w = img.shape[:2]
    horizon_frac = settings.HORIZON_FRAC if horizon_frac is None else horizon_frac

    if ml_available() and settings.SAM_CHECKPOINT:
        m = _sam_floor_mask(img)
        if m is not None:
            return m
    return geometric_floor_mask(h, w, horizon_frac)


def geometric_floor_mask(h: int, w: int, horizon_frac: float) -> np.ndarray:
    """Floor = trapezoid below the horizon, narrowing toward the back wall."""
    horizon_y = int(h * horizon_frac)
    poly = np.array(
        [
            [0, h - 1],
            [w - 1, h - 1],
            [int(w * 0.85), horizon_y],
            [int(w * 0.15), horizon_y],
        ],
        dtype=np.int32,
    )
    mask = np.zeros((h, w), dtype=np.uint8)
    cv2.fillPoly(mask, [poly], 255)
    return mask


def _sam_floor_mask(img: np.ndarray):
    """Point-prompted SAM floor isolation. Returns None on any problem."""
    ckpt = Path(settings.SAM_CHECKPOINT)
    if not ckpt.exists():
        return None
    try:
        import torch
        from segment_anything import SamPredictor, sam_model_registry

        sam = sam_model_registry[settings.SAM_MODEL_TYPE](checkpoint=str(ckpt))
        sam.to("cuda" if torch.cuda.is_available() else "cpu")
        predictor = SamPredictor(sam)
        rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        predictor.set_image(rgb)
        h, w = img.shape[:2]
        # Prompt with points along the bottom-center, where floor is most likely.
        pts = np.array(
            [[w // 2, int(h * 0.93)], [w // 3, int(h * 0.88)], [2 * w // 3, int(h * 0.88)]]
        )
        labels = np.ones(len(pts), dtype=np.int32)
        masks, scores, _ = predictor.predict(
            point_coords=pts, point_labels=labels, multimask_output=True
        )
        best = masks[int(np.argmax(scores))]
        return (best.astype("uint8") * 255)
    except Exception:
        return None