import json
import builtins
import sys
from pathlib import Path
from types import SimpleNamespace
from unittest.mock import patch

import numpy as np

SOURCE_ROOT = Path(__file__).parents[1] / "src"
sys.path.insert(0, str(SOURCE_ROOT))

from anyimage.common.selection import (
    SelectionMask,
    SelectionPath,
    rasterize_selection_path,
)
from anyimage.common.depth import (
    depth_uniform_scale,
    load_depth_metadata,
)
from anyimage.operators.cutout_tool import geometry as cutout_geometry


from anyimage.operators.cutout_tool import object as cutout_object
from anyimage.operators.cutout_tool import shape as cutout_shape
from anyimage.operators.cutout_tool.geometry import (
    SelectionContours,
    extract_selection_contours,
    poisson_balloon_profile,
)
from anyimage.server.models import moge2 as moge2_model
from anyimage.server.models import da3 as da3_model
from anyimage.server.models.geometry import GeometryFrame


def test_cutout_mesh_detail_uses_the_cropped_analysis_scale():
    assert cutout_geometry.limit_cutout_mesh_resolution(256, (4096, 1024), 2048) == 128
    assert cutout_geometry.limit_cutout_mesh_resolution(256, (1024, 512), 2048) == 256
    assert cutout_geometry.limit_cutout_mesh_resolution(8, (8192, 4096), 2048) == 4


def test_depth_metadata_loader_uses_one_schema_for_plane_and_cutout(tmp_path):
    path = tmp_path / "depth.json"
    path.write_text(
        json.dumps(
            {
                "image_size": [3, 2],
                "intrinsics": np.eye(3).tolist(),
                "reference_depth": 1.0,
            }
        ),
        encoding="utf-8",
    )

    metadata = load_depth_metadata(path)

    assert metadata["image_size"] == (3, 2)
    assert metadata["intrinsics"] == ((1.0, 0.0, 0.0), (0.0, 1.0, 0.0), (0.0, 0.0, 1.0))
    assert metadata["reference_depth"] == 1.0


def test_cutout_node_groups_are_loaded_from_the_asset():
    loaded = object()
    with patch.object(
        cutout_object, "load_node_group", return_value=loaded
    ) as load:
        for shape in cutout_shape.CUTOUT_NODE_GROUP_NAMES:
            assert cutout_object.load_cutout_node_group() is loaded
            load.assert_called_with(cutout_shape.CUTOUT_NODE_GROUP_NAME)
        assert load.call_count == 4


def test_base_plane_fit_is_stable_for_a_slender_diagonal_domain():
    along = np.linspace(-1.0, 1.0, 101)
    across = np.resize(np.asarray((-0.01, 0.0, 0.01)), len(along))
    coordinates = np.column_stack(
        (
            along - across,
            along + across,
        )
    )
    values = 0.4 * along + np.sign(across) * 0.05

    plane = cutout_geometry._stable_plane(coordinates, values)
    selected_plane = coordinates @ plane[:2] + plane[2]
    corners = np.asarray(((-1.0, -1.0), (-1.0, 1.0), (1.0, -1.0), (1.0, 1.0)))
    corner_plane = corners @ plane[:2] + plane[2]

    assert np.ptp(corner_plane) <= np.ptp(selected_plane) * 1.25
    assert np.corrcoef(selected_plane, along)[0, 1] > 0.99


def test_depth_balloon_plane_calibration_matches_cutout_axes():
    vertices = []
    uv = []
    camera_depth = np.empty((3, 3), dtype=np.float32)
    for row in range(3):
        local_z = 1.0 - row
        for column in range(3):
            local_y = column - 1.0
            vertices.append((0.0, local_y, local_z))
            uv.append((column * 0.5, 1.0 - row * 0.5))
            camera_depth[row, column] = 2.0 + 0.2 * local_y + 0.1 * local_z

    reference, direction = cutout_geometry.depth_balloon_projection_parameters(
        vertices,
        uv,
        camera_depth,
        np.ones((3, 3), dtype=bool),
        0.5,
    )
    expected_direction = np.asarray((-0.1, 0.05, 1.0), dtype=np.float64)
    expected_direction /= np.linalg.norm(expected_direction)

    assert np.isclose(
        reference,
        np.percentile(camera_depth, 95.0) * 0.5,
        atol=1e-6,
    )
    assert np.allclose(direction, expected_direction, atol=1e-6)


def test_depth_balloon_reference_uses_the_shared_upper_depth_percentile():
    size = 10
    row, column = np.indices((size, size))
    camera_depth = np.full((size, size), 2.0, dtype=np.float32)
    camera_depth[(row + column) % 2 == 0] = 2.2
    vertices = []
    uv = []
    for y in range(size):
        for x in range(size):
            vertices.append((0.0, float(x), float(size - 1 - y)))
            uv.append((x / (size - 1), 1.0 - y / (size - 1)))

    reference, direction = cutout_geometry.depth_balloon_projection_parameters(
        vertices,
        uv,
        camera_depth,
        np.ones((size, size), dtype=bool),
        1.0,
    )

    assert reference > float(np.median(camera_depth))
    assert np.isclose(reference, 2.2, atol=1e-6)
    assert np.allclose(direction, (0.0, 0.0, 1.0), atol=1e-6)


def mesh_component_count(faces):
    adjacency = {}
    for face in faces:
        for vertex in face:
            adjacency.setdefault(int(vertex), set())
        for start, end in zip(face, np.roll(face, -1)):
            adjacency[int(start)].add(int(end))
            adjacency[int(end)].add(int(start))
    remaining = set(adjacency)
    count = 0
    while remaining:
        count += 1
        pending = [remaining.pop()]
        while pending:
            vertex = pending.pop()
            connected = adjacency[vertex] & remaining
            remaining.difference_update(connected)
            pending.extend(connected)
    return count


def test_depth_surface_projection_uses_one_scale_for_all_camera_axes():
    points = np.zeros((51, 101, 3), dtype=np.float32)
    points[..., 2] = np.linspace(1.0, 3.0, 51 * 101).reshape(51, 101)
    metadata = {
        "intrinsics": (
            (100.0, 0.0, 50.0),
            (0.0, 100.0, 25.0),
            (0.0, 0.0, 1.0),
        ),
        "image_size": (101, 51),
        "reference_depth": 2.0,
    }

    uniform_scale = depth_uniform_scale(
        SimpleNamespace(),
        metadata,
        bounds=(-2.0, 2.0, -1.0, 1.0),
    )

    assert uniform_scale == 2.0
    reference_depth = cutout_geometry.metric_depth_reference(
        points[..., 2],
        np.ones((51, 101), dtype=bool),
        uniform_scale,
    )
    assert np.isclose(reference_depth, 5.8, atol=1e-6)


def test_geometry_frame_requires_matching_float32_fields():
    with np.testing.assert_raises_regex(ValueError, "validity"):
        GeometryFrame(
            depth=np.ones((2, 2), dtype=np.float32),
            validity=np.ones((1, 2), dtype=np.float32),
            intrinsics=np.eye(3, dtype=np.float32),
        )
    with np.testing.assert_raises_regex(ValueError, "normal"):
        GeometryFrame(
            depth=np.ones((2, 2), dtype=np.float32),
            validity=np.ones((2, 2), dtype=np.float32),
            intrinsics=np.eye(3, dtype=np.float32),
            normal=np.ones((2, 2, 2), dtype=np.float32),
        )


def test_moge_infer_returns_geometry_frames_with_model_validity():
    mask = np.asarray(((True, False), (True, True)))
    prediction = {
        "depth": np.ones((2, 2), dtype=np.float32),
        "normal": np.zeros((2, 2, 3), dtype=np.float32),
        "mask": mask,
        "intrinsics": np.asarray(
            ((1.0, 0.0, 0.5), (0.0, 1.0, 0.5), (0.0, 0.0, 1.0)),
            dtype=np.float32,
        ),
        "points": np.ones((2, 2, 3), dtype=np.float32),
    }

    with patch.object(moge2_model, "infer_one", return_value=prediction):
        frames = moge2_model.infer(None, {}, ["image"], include_points=True)

    assert len(frames) == 1
    frame = frames[0]
    assert np.array_equal(frame.validity, mask.astype(np.float32))
    assert np.array_equal(frame.normal, prediction["normal"])
    assert np.array_equal(frame.points, prediction["points"])
    assert np.array_equal(
        frame.intrinsics,
        np.asarray(((2.0, 0.0, 1.0), (0.0, 2.0, 1.0), (0.0, 0.0, 1.0))),
    )


def test_moge_z_depth_frame_does_not_require_points():
    prediction = {
        "depth": np.ones((2, 2), dtype=np.float32),
        "normal": np.zeros((2, 2, 3), dtype=np.float32),
        "mask": np.ones((2, 2), dtype=bool),
        "intrinsics": np.eye(3, dtype=np.float32),
        "points": np.ones((2, 2, 3), dtype=np.float32),
    }

    with patch.object(moge2_model, "infer_one", return_value=prediction):
        frame = moge2_model.infer(None, {}, ["image"], include_points=False)[0]

    assert frame.points is None
    assert frame.validity.all()


def test_da3_geometry_frames_use_confidence_as_validity_and_derive_points():
    prediction = SimpleNamespace(
        depth=np.ones((2, 2, 3), dtype=np.float32),
        conf=np.asarray(
            (
                ((0.25, 0.5, 0.75), (1.0, 0.5, 0.25)),
                ((0.1, 0.2, 0.3), (0.4, 0.5, 0.6)),
            ),
            dtype=np.float32,
        ),
        intrinsics=np.repeat(np.eye(3, dtype=np.float32)[None], 2, axis=0),
    )

    frames = da3_model.geometry_frames(prediction)

    assert len(frames) == 2
    assert frames[0].normal is None
    assert np.array_equal(frames[1].validity, prediction.conf[1])
    assert frames[0].points.shape == (2, 3, 3)
    assert frames[0].points.dtype == np.float32


def test_selection_contours_group_holes_with_their_connected_component():
    alpha = np.zeros((120, 200), dtype=np.uint8)
    alpha[20:100, 50:160] = 255
    alpha[45:75, 80:120] = 0
    alpha[5:15, 5:15] = 255

    selection_contours = extract_selection_contours(
        SelectionMask(alpha, (0, 0, 200, 120)),
        max_resolution=200,
    )

    assert selection_contours.image_size == (200, 120)
    assert sorted(len(component) for component in selection_contours.components) == [1, 2]
    contours = [
        contour for component in selection_contours.components for contour in component
    ]
    bounds = [
        (*np.min(contour, axis=0), *np.max(contour, axis=0)) for contour in contours
    ]
    assert (50.0, 20.0, 160.0, 100.0) in bounds
    assert (80.0, 45.0, 120.0, 75.0) in bounds
    assert (5.0, 5.0, 15.0, 15.0) in bounds


def test_selection_contours_accept_blender_float_mask_with_a_hole():
    alpha = np.zeros((120, 200), dtype=np.float32)
    alpha[20:100, 50:160] = 1.0
    alpha[45:75, 80:120] = 0.0

    selection_contours = extract_selection_contours(
        SelectionMask(alpha, (0, 0, 200, 120)),
        max_resolution=200,
    )

    assert selection_contours.image_size == (200, 120)
    assert len(selection_contours.components) == 1
    assert len(selection_contours.components[0]) == 2


def test_cutout_crop_bounds_map_to_the_source_plane():
    source = SimpleNamespace(
        data=SimpleNamespace(size=(200, 100)),
        empty_display_size=2.0,
        empty_image_offset=(-0.5, -0.5),
    )

    bounds = cutout_geometry.crop_plane_bounds(
        source,
        (50, 25, 150, 75),
        (200, 100),
    )

    np.testing.assert_allclose(bounds, (-0.5, 0.5, -0.25, 0.25))


def test_depth_surface_cutout_uses_half_alpha_boundary():
    alpha = np.zeros((12, 16), dtype=np.float32)
    alpha[2:10, 2:14] = 0.49
    alpha[4:8, 5:11] = 0.51

    selection_contours = cutout_geometry.extract_cutout_contours(
        SelectionMask(alpha, (0, 0, 16, 12)),
        64,
        threshold=0.5,
    )
    contour = np.asarray(selection_contours.components[0][0])

    assert selection_contours.image_size == (16, 12)
    np.testing.assert_allclose(contour.min(axis=0), (5.0, 4.0))
    np.testing.assert_allclose(contour.max(axis=0), (11.0, 8.0))


def test_blender_cutout_alpha_pipeline_does_not_import_pillow():
    original_import = builtins.__import__

    def without_pillow(name, *args, **kwargs):
        if name == "PIL" or name.startswith("PIL."):
            raise ModuleNotFoundError("Pillow is unavailable in Blender")
        return original_import(name, *args, **kwargs)

    with patch("builtins.__import__", side_effect=without_pillow):
        selection_mask = rasterize_selection_path(
            (48, 32),
            SelectionPath(points=((6, 4), (42, 4), (42, 28), (6, 28))),
            antialias=True,
        )
        antialiased = selection_mask.full_values((48, 32))
        selection_contours = extract_selection_contours(
            SelectionMask(antialiased, (0, 0, 48, 32)),
            max_resolution=24,
        )

    assert selection_contours.image_size == (48, 32)
    assert selection_contours.components
    assert np.any((antialiased > 0.0) & (antialiased < 1.0))


def test_blender_cutout_contour_resolution_tracks_mesh_detail():
    expected = SelectionContours(
        components=(((((0, 0), (1, 0), (0, 1))),),),
        image_size=(1, 1),
    )
    selection_mask = SelectionMask(np.ones((2, 2), dtype=np.float32), (0, 0, 2, 2))
    with patch.object(cutout_geometry, "extract_selection_contours", return_value=expected) as trace:
        assert cutout_geometry.extract_cutout_contours(selection_mask, 32) == expected
        assert cutout_geometry.extract_cutout_contours(selection_mask, 256) == expected

    assert trace.call_args_list[0].kwargs["max_resolution"] == 512
    assert trace.call_args_list[1].kwargs["max_resolution"] == 1024


def test_cutout_sampling_spacing_uses_the_complete_source_plane():
    source = SimpleNamespace(matrix_world=object())
    with (
        patch.object(cutout_geometry, "image_empty_bounds", return_value=(-2.0, 2.0, -1.0, 1.0)),
        patch.object(
            cutout_geometry,
            "world_plane_metric",
            return_value=(((1.0, 0.0), (0.0, 1.0)), 1.0),
        ),
    ):
        spacing = cutout_geometry.cutout_sampling_spacing(source, 200)

    assert spacing == 0.02


def test_cutout_cleanup_value_scales_the_ui_without_changing_node_area():
    assert np.isclose(
        cutout_geometry.cutout_cleanup_value(
            0.25,
            ((1.0, 0.0), (0.0, 1.0)),
        ),
        6.25,
    )
    assert np.isclose(
        cutout_geometry.cutout_cleanup_value(
            0.25,
            ((4.0, 0.0), (0.0, 9.0)),
        ),
        6.25 / 6.0,
    )


def test_cutout_depth_calibration_uses_five_percent_transparency_threshold():
    assert not cutout_geometry.use_adaptive_depth_calibration(0.0499)
    assert cutout_geometry.use_adaptive_depth_calibration(0.05)


def test_poisson_balloon_profile_is_zero_on_boundary_and_positive_inside():
    points = np.asarray(
        ((0, 0), (2, 0), (2, 2), (0, 2), (1, 1)),
        dtype=float,
    )
    faces = ((0, 1, 4), (1, 2, 4), (2, 3, 4), (3, 0, 4))
    boundary = np.asarray((True, True, True, True, False))

    profile = poisson_balloon_profile(points, faces, boundary)

    assert np.all(profile[:4] == 0.0)
    assert profile[4] > 0.0


def test_poisson_balloon_profile_scales_in_world_units():
    points = np.asarray(
        ((0, 0), (2, 0), (2, 2), (0, 2), (1, 1)),
        dtype=float,
    )
    faces = ((0, 1, 4), (1, 2, 4), (2, 3, 4), (3, 0, 4))
    boundary = np.asarray((True, True, True, True, False))

    original = poisson_balloon_profile(points, faces, boundary)
    scaled = poisson_balloon_profile(points * 2.0, faces, boundary)

    assert np.isclose(scaled[4], original[4] * 2.0, atol=1e-6)


def test_poisson_balloon_profile_makes_wide_selections_thicker():
    faces = ((0, 1, 4), (1, 2, 4), (2, 3, 4), (3, 0, 4))
    boundary = np.asarray((True, True, True, True, False))
    wide = np.asarray(
        ((0, 0), (4, 0), (4, 2), (0, 2), (2, 1)),
        dtype=float,
    )
    narrow = np.asarray(
        ((0, 0), (4, 0), (4, 0.6), (0, 0.6), (2, 0.3)),
        dtype=float,
    )

    wide_profile = poisson_balloon_profile(wide, faces, boundary)
    narrow_profile = poisson_balloon_profile(narrow, faces, boundary)

    assert wide_profile[4] > narrow_profile[4]


def test_teddy_balloon_profile_uses_chordal_axis_radius():
    points = np.asarray(
        ((0, 0), (4, 0), (4, 2), (0, 2), (2, 1)),
        dtype=float,
    )
    faces = ((0, 1, 4), (1, 2, 4), (2, 3, 4), (3, 0, 4))
    boundary = np.asarray((True, True, True, True, False))
    nodes = np.asarray(((1, 1), (3, 1)), dtype=float)

    with patch.object(
        cutout_geometry,
        "chordal_axis",
        return_value=(nodes, [(0, 1)], np.asarray((1.0, 1.0))),
    ):
        profile = cutout_geometry.teddy_balloon_heights(
            points,
            faces,
            points[:4],
            boundary,
            0.45,
        )

    assert np.all(profile[:4] == 0.0)
    assert np.isclose(profile[4], 0.45, atol=1e-6)


def test_cutout_domain_adapts_boundary_spacing_for_a_small_valid_selection():
    components = ((((2046, 2046), (2050, 2046), (2050, 2050), (2046, 2050)),),)

    domain = cutout_geometry.cutout_domain(
        SelectionContours(components, (4096, 4096)),
        (-1.0, 1.0, -1.0, 1.0),
        0.5,
        ((1.0, 0.0), (0.0, 1.0)),
    )

    assert len(domain[0]) >= 3
    assert len(domain[1]) >= 1


def test_cutout_domain_builds_a_full_image_at_minimum_detail():
    components = ((((0, 0), (4096, 0), (4096, 4096), (0, 4096)),),)

    domain = cutout_geometry.cutout_domain(
        SelectionContours(components, (4096, 4096)),
        (-1.0, 1.0, -1.0, 1.0),
        0.5,
        ((1.0, 0.0), (0.0, 1.0)),
    )

    assert len(domain[0]) >= 4
    assert len(domain[1]) >= 2


def test_cutout_domain_centers_vertices_on_crop_bounds():
    components = ((((0, 0), (100, 0), (100, 50), (0, 50)),),)

    vertices, _faces, uv, _boundary, _planar, _profile = cutout_geometry.cutout_domain(
        SelectionContours(components, (100, 50)),
        (2.0, 6.0, -3.0, 1.0),
        0.5,
        ((1.0, 0.0), (0.0, 1.0)),
    )

    center = (vertices[:, :2].min(axis=0) + vertices[:, :2].max(axis=0)) * 0.5
    assert np.allclose(center, 0.0, atol=1e-6)
    expected_uv = vertices[:, :2] / np.asarray((4.0, 4.0)) + 0.5
    assert np.allclose(uv, expected_uv, atol=1e-6)


def test_cutout_domain_keeps_overlapping_selection_components_disconnected():
    components = (
        (
            ((20, 0), (100, 0), (100, 80), (20, 80)),
            ((50, 20), (50, 60), (70, 60), (70, 20)),
        ),
        (((0, 20), (40, 20), (40, 50), (0, 50)),),
    )
    domain = cutout_geometry.cutout_domain(
        SelectionContours(components, (120, 80)),
        (0.0, 1.0, 0.0, 1.0),
        0.025,
        ((1.0, 0.0), (0.0, 1.0)),
    )

    centers = domain[2][domain[1]].mean(axis=1)
    in_hole = (
        (centers[:, 0] > 0.45)
        & (centers[:, 0] < 0.55)
        & (centers[:, 1] > 0.35)
        & (centers[:, 1] < 0.65)
    )
    in_left_component = (
        (centers[:, 0] > 0.02)
        & (centers[:, 0] < 0.12)
        & (centers[:, 1] > 0.4)
        & (centers[:, 1] < 0.75)
    )

    assert not in_hole.any()
    assert in_left_component.any()
    assert mesh_component_count(domain[1]) == 2


def test_selection_calibration_uses_the_final_selection_mask():
    selection_mask = SelectionMask(
        np.asarray(((1.0, 1.0), (1.0, 0.0)), dtype=np.float32),
        (0, 0, 2, 2),
    )

    valid, transparent_fraction = cutout_geometry.selection_calibration(selection_mask, (2, 2))

    assert np.array_equal(valid, ((True, True), (True, False)))
    assert transparent_fraction == 0.25