import os
import sys
from types import SimpleNamespace

import numpy as np
import pytest
from PIL import Image


PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
if PROJECT_ROOT not in sys.path:
    sys.path.insert(0, PROJECT_ROOT)


import tools.model_view_tools as model_view_tools_module
import tools.reference_image as reference_image_module
from tools.model_view_tools import get_visual_feedback


@pytest.mark.asyncio
async def test_get_visual_feedback_uses_latest_uploaded_image_for_both_llm_steps(
    monkeypatch, tmp_path
):
    image_path = tmp_path / "query_image_001.png"
    image_path.write_bytes(b"fake-png-bytes")
    render_path = tmp_path / "render.png"
    render_path.write_bytes(b"fake-render")

    class FakeContext:
        def retrieve_full_messages(self):
            return [
                SimpleNamespace(
                    role="user",
                    content=[
                        {"type": "text", "text": "Inspect this part"},
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": "data:image/png;base64,abcd",
                                "local_path": str(image_path),
                            },
                        },
                    ],
                )
            ]

    monkeypatch.setattr(
        reference_image_module,
        "get_current_context",
        lambda: FakeContext(),
    )
    monkeypatch.setattr(
        model_view_tools_module, "_require_simplecad_renderer", lambda: None
    )
    monkeypatch.setattr(
        model_view_tools_module, "print_tool_output", lambda *args, **kwargs: None
    )

    class FakeSketchPad:
        async def set_item(self, key, value, ttl=None, summary=None, tags=None):
            return key

    monkeypatch.setattr(
        model_view_tools_module,
        "get_current_sketch_pad",
        lambda: FakeSketchPad(),
    )

    captured = {}

    async def fake_question_generator(user_query, code, query_image_path):
        captured["question_image_path"] = (
            str(query_image_path.path) if query_image_path else None
        )
        return "Checklist"

    async def fake_visual_feedback_generator(
        questions, multi_view_results, query_image_path
    ):
        captured["visual_image_path"] = (
            str(query_image_path.path) if query_image_path else None
        )
        captured["multi_view_results"] = str(multi_view_results.path)
        return "Looks correct\nPASS"

    monkeypatch.setattr(
        model_view_tools_module,
        "question_generator",
        fake_question_generator,
    )
    monkeypatch.setattr(
        model_view_tools_module,
        "visual_feedback_generator",
        fake_visual_feedback_generator,
    )

    from SimpleLLMFunc.type import ImgPath

    monkeypatch.setattr(
        model_view_tools_module,
        "render_multi_view_model",
        lambda model_path, output_path: ImgPath(render_path, detail="high"),
    )

    result = await get_visual_feedback(
        user_query="Inspect this part",
        code="result = None",
        model_path="./part/model.stl",
    )

    expected_image_path = str(image_path.resolve())
    assert captured["question_image_path"] == expected_image_path
    assert captured["visual_image_path"] == expected_image_path
    assert captured["multi_view_results"] == str(render_path.resolve())
    assert "Model path: ./part/model.stl" in result


def test_camera_relative_light_rig_tracks_camera_direction() -> None:
    for view_dir in (
        np.array([1.0, 0.0, 0.0]),
        np.array([-1.0, 0.0, 0.0]),
        np.array([0.0, 0.0, 1.0]),
        np.array([1.0, 1.0, -1.0]),
    ):
        normalized_view = view_dir / np.linalg.norm(view_dir)
        light_dirs, light_weights, ambient = (
            model_view_tools_module._camera_relative_light_rig(normalized_view)
        )

        assert ambient > 0.0
        assert len(light_dirs) == len(light_weights) >= 3
        assert np.dot(light_dirs[0], normalized_view) > 0.85
        assert all(
            abs(np.linalg.norm(light_dir) - 1.0) < 1e-6 for light_dir in light_dirs
        )


def test_camera_relative_shading_keeps_front_faces_bright_across_views() -> None:
    base_color = np.array([0.72, 0.76, 0.81], dtype=float)
    front_face_brightness: list[float] = []

    for view_dir in (
        np.array([1.0, 0.0, 0.0]),
        np.array([-1.0, 0.0, 0.0]),
        np.array([0.0, 1.0, 0.0]),
        np.array([1.0, 1.0, -1.0]),
    ):
        normalized_view = view_dir / np.linalg.norm(view_dir)
        normals = np.stack([normalized_view, -normalized_view])
        shaded = model_view_tools_module._shade_normals_camera_relative(
            normals,
            base_color,
            normalized_view,
        )
        front_face_brightness.append(float(np.mean(shaded[0, :3])))
        assert float(np.mean(shaded[0, :3])) > float(np.mean(shaded[1, :3]))

    assert min(front_face_brightness) > 0.55
    assert max(front_face_brightness) - min(front_face_brightness) < 0.12


def test_surface_shading_normals_keep_planar_triangles_consistent() -> None:
    class FakeVector:
        def __init__(self, x: float, y: float, z: float) -> None:
            self.x = x
            self.y = y
            self.z = z

    class FakeCadFace:
        def geomType(self) -> str:
            return "PLANE"

        def normalAt(self, location=None):
            return FakeVector(0.0, 0.0, 1.0)

    class FakeFace:
        cq_face = FakeCadFace()

    tri_pts = np.array(
        [
            [[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [1.0, 1.0, 0.0]],
            [[0.0, 0.0, 0.0], [1.0, 1.0, 0.0], [0.0, 1.0, 0.0]],
        ],
        dtype=float,
    )

    normals = model_view_tools_module._compute_surface_shading_normals(
        FakeFace(), tri_pts
    )

    assert normals.shape == (2, 3)
    assert np.allclose(normals[0], [0.0, 0.0, 1.0])
    assert np.allclose(normals[1], [0.0, 0.0, 1.0])


def test_direct_cad_renderer_is_used_for_brep_like_formats() -> None:
    assert model_view_tools_module._should_use_direct_cad_renderer("part.step") is True
    assert model_view_tools_module._should_use_direct_cad_renderer("part.stp") is True
    assert model_view_tools_module._should_use_direct_cad_renderer("part.brep") is True
    assert model_view_tools_module._should_use_direct_cad_renderer("part.bin") is True
    assert model_view_tools_module._should_use_direct_cad_renderer("part.stl") is False


def test_prefer_cad_native_model_path_uses_step_over_stl(tmp_path) -> None:
    stl_path = tmp_path / "model.stl"
    step_path = tmp_path / "model.step"
    stl_path.write_text("solid", encoding="utf-8")
    step_path.write_text("step", encoding="utf-8")

    selected = model_view_tools_module._prefer_cad_native_model_path(str(stl_path))

    assert selected == str(step_path.resolve())


def test_prefer_cad_native_model_path_keeps_stl_without_step(tmp_path) -> None:
    stl_path = tmp_path / "model.stl"
    stl_path.write_text("solid", encoding="utf-8")

    selected = model_view_tools_module._prefer_cad_native_model_path(str(stl_path))

    assert selected == str(stl_path.resolve())


def test_load_renderable_shapes_flattens_step_compounds(monkeypatch) -> None:
    class FakeCadShape:
        def __init__(self, shape_type: str, solids=None) -> None:
            self._shape_type = shape_type
            self._solids = list(solids or [])

        def ShapeType(self) -> str:
            return self._shape_type

        def Solids(self):
            return list(self._solids)

    class FakeWrappedSolid:
        def __init__(self, obj) -> None:
            self.obj = obj

    class FakeWorkplane:
        def __init__(self, values) -> None:
            self._values = values

        def vals(self):
            return list(self._values)

    solid_a = FakeCadShape("Solid")
    solid_b = FakeCadShape("Solid")
    compound = FakeCadShape("Compound", solids=[solid_a, solid_b])

    fake_cq = SimpleNamespace(
        importers=SimpleNamespace(
            importShape=lambda import_type, model_path: FakeWorkplane([compound])
        )
    )

    monkeypatch.setattr(model_view_tools_module, "cq", fake_cq)
    monkeypatch.setattr(model_view_tools_module, "ScadSolid", FakeWrappedSolid)
    monkeypatch.setattr(
        model_view_tools_module, "_require_simplecad_renderer", lambda: None
    )

    result = model_view_tools_module._load_renderable_shapes("part.step")

    assert [wrapped.obj for wrapped in result] == [solid_a, solid_b]


def test_feature_edge_mask_detects_normal_and_depth_discontinuities() -> None:
    mask = np.ones((8, 8), dtype=bool)
    depth = np.zeros((8, 8), dtype=float)
    normals = np.zeros((8, 8, 3), dtype=float)
    normals[:, :4] = np.array([0.0, 0.0, 1.0])
    normals[:, 4:] = np.array([1.0, 0.0, 0.0])

    edge_mask = model_view_tools_module._compute_feature_edge_mask(
        mask,
        depth,
        normals,
        depth_jump_threshold=10.0,
        normal_cos_threshold=0.95,
    )

    assert edge_mask[:, 3:5].any()

    depth[:, 4:] = 4.0
    normals[:, :] = np.array([0.0, 0.0, 1.0])
    edge_mask = model_view_tools_module._compute_feature_edge_mask(
        mask,
        depth,
        normals,
        depth_jump_threshold=1.0,
        normal_cos_threshold=0.95,
    )

    assert edge_mask[:, 3:5].any()


def test_direct_rasterizer_supersamples_and_preserves_target_size() -> None:
    triangles = [
        np.array([[0.0, 0.0, 1.0], [1.0, 0.0, 1.0], [1.0, 1.0, 1.0]], dtype=float),
        np.array([[0.0, 0.0, 1.0], [1.0, 1.0, 1.0], [0.0, 1.0, 1.0]], dtype=float),
    ]
    normals = [
        np.array([0.0, 0.0, 1.0], dtype=float),
        np.array([0.0, 0.0, 1.0], dtype=float),
    ]

    image = model_view_tools_module._rasterize_projected_triangles(
        triangles,
        normals,
        image_size=(48, 48),
        zoom=4.0,
        background_rgb=np.array([255, 255, 255], dtype=np.uint8),
        fill_rgb=np.array([180, 190, 200], dtype=np.uint8),
        outline_rgb=np.array([0, 0, 0], dtype=np.uint8),
    )

    pixels = np.asarray(image)
    unique_colors = np.unique(pixels.reshape(-1, 3), axis=0)

    assert image.size == (48, 48)
    assert len(unique_colors) > 3


def test_axis_triad_overlay_draws_small_corner_marker() -> None:
    image = Image.new("RGB", (320, 320), "white")
    annotated = model_view_tools_module._add_axis_triad_overlay(
        image,
        np.array([1.0, 1.0, 1.0]) / np.sqrt(3.0),
    )

    original = np.asarray(image)
    updated = np.asarray(annotated)
    diff = np.abs(updated.astype(int) - original.astype(int)).sum(axis=2)
    changed_pixels = np.argwhere(diff > 0)

    assert changed_pixels.size > 0
    assert int(changed_pixels[:, 0].max()) > 220
    assert int(changed_pixels[:, 1].max()) < 120