"""Tests for the segment_with_sam21/points_via_sam21 decomposition.

We exercise the path that the find_model endpoint will take:
  segment_with_sam21(nuwa, sam21, bbox, sw) -> Sam21Segmentation
and the legacy wrapper still works.
"""
from __future__ import annotations

from unittest.mock import MagicMock

import cv2
import numpy as np
import pytest

from presa_mano.perception.sam21_bridge import (
    Sam21Segmentation,
    segment_with_sam21,
    points_via_sam21,
    Sam21Fallback,
)
from presa_mano.timing import Stopwatch


def _fake_nuwa():
    """Returns a MagicMock Nuwa that yields a 4x4 black RGB jpeg and a 4x4
    uint16 depth PNG (all 1000 mm) + a minimal Nuwa intrinsics dict."""
    rgb = np.zeros((4, 4, 3), dtype=np.uint8)
    ok, jpg = cv2.imencode(".jpg", rgb)
    assert ok
    depth = np.full((4, 4), 1000, dtype=np.uint16)
    ok, dpng = cv2.imencode(".png", depth)
    assert ok
    intr = {
        "rgb": {"fx": 600.0, "fy": 600.0, "cx": 2.0, "cy": 2.0},
        "ir":  {"fx": 600.0, "fy": 600.0, "cx": 2.0, "cy": 2.0,
                "extrinsics": {"R": [[1,0,0],[0,1,0],[0,0,1]],
                               "t": [0.0, 0.0, 0.0]}},
    }
    n = MagicMock()
    n.rgb_jpeg.return_value = bytes(jpg)
    n.depth_png.return_value = bytes(dpng)
    n.intrinsics.return_value = intr
    return n, bytes(jpg)


def _fake_sam21(points_m: np.ndarray):
    s = MagicMock()
    s.segment_pointcloud.return_value = points_m
    return s


def test_segment_with_sam21_returns_segmentation(monkeypatch):
    nuwa, rgb_jpeg = _fake_nuwa()
    pts_m = np.array([[0.01, 0.02, 0.5], [0.011, 0.022, 0.49]])
    sam21 = _fake_sam21(pts_m)
    # Niente plane filter qui — segment_with_sam21 NON deve applicarlo
    monkeypatch.setattr(
        "presa_mano.perception.sam21_bridge.get_work_plane", lambda: None
    )
    sw = Stopwatch()
    seg = segment_with_sam21(nuwa, sam21, bbox=(0, 0, 4, 4), sw=sw)
    assert isinstance(seg, Sam21Segmentation)
    np.testing.assert_array_equal(seg.points_cv_m, pts_m)
    assert seg.rgb_jpeg == rgb_jpeg
    assert "width" in seg.intrinsics_rgb and "height" in seg.intrinsics_rgb


def test_points_via_sam21_still_applies_plane_filter(monkeypatch):
    """L'API legacy non deve cambiare: ritorna (pts, plane_info) e applica
    apply_below_plane_filter come prima."""
    nuwa, _rgb = _fake_nuwa()
    pts_m = np.array([[0.0, 0.0, 0.5]] * 20)
    sam21 = _fake_sam21(pts_m)
    monkeypatch.setattr(
        "presa_mano.perception.sam21_bridge.get_work_plane", lambda: None
    )
    sw = Stopwatch()
    pts, info = points_via_sam21(nuwa, sam21, bbox=(0, 0, 4, 4), sw=sw)
    assert pts.shape == (20, 3)
    assert info["plane_applied"] is False
    assert info["n_points_before_plane"] == 20


def test_points_via_sam21_too_few_points_raises_fallback(monkeypatch):
    nuwa, _rgb = _fake_nuwa()
    pts_m = np.array([[0.0, 0.0, 0.5]] * 5)  # <10
    sam21 = _fake_sam21(pts_m)
    monkeypatch.setattr(
        "presa_mano.perception.sam21_bridge.get_work_plane", lambda: None
    )
    sw = Stopwatch()
    with pytest.raises(Sam21Fallback):
        points_via_sam21(nuwa, sam21, bbox=(0, 0, 4, 4), sw=sw)