"""Streamable HTTP guided-surface regressions for default spatial support."""

from __future__ import annotations

import asyncio
import textwrap
from pathlib import Path

import pytest
from fastmcp.exceptions import ToolError

from ._guided_surface_harness import result_payload, run_streamable_server, streamable_client, write_server_script

_PATCHED_GUIDED_STREAMABLE_SERVER = textwrap.dedent(
    """
    from server.adapters.mcp.areas import router as router_area
    import server.adapters.mcp.areas.collection as collection_area
    import server.adapters.mcp.areas.mesh as mesh_area
    import server.adapters.mcp.areas.modeling as modeling_area
    import server.adapters.mcp.areas.scene as scene_area
    import server.adapters.mcp.router_helper as router_helper
    import server.infrastructure.di as di


    class RouterHandler:
        def set_goal(self, goal, resolved_params=None):
            return {
                "status": "no_match",
                "continuation_mode": "guided_manual_build",
                "workflow": None,
                "resolved": {},
                "unresolved": [],
                "resolution_sources": {},
                "phase_hint": "build",
                "message": "Continue on the guided build surface.",
            }

        def clear_goal(self):
            return "cleared"


    class SceneHandler:
        def list_objects(self):
            return [
                {"name": "Squirrel_Body", "type": "MESH"},
                {"name": "Squirrel_Head", "type": "MESH"},
                {"name": "ForeL", "type": "MESH"},
            ]

        def clean_scene(self, keep_lights_and_cameras):
            return "Scene cleaned."

        def get_scope_graph(self, target_object=None, target_objects=None, collection_name=None):
            names = [name for name in [target_object, *(target_objects or [])] if name]
            primary = target_object or (names[0] if names else None)
            return {
                "scope_kind": "object_set" if len(names) > 1 else "single_object",
                "primary_target": primary,
                "object_names": names,
                "object_count": len(names),
                "object_roles": [],
            }

        def get_relation_graph(
            self,
            target_object=None,
            target_objects=None,
            collection_name=None,
            goal_hint=None,
            include_truth_payloads=False,
        ):
            names = [name for name in [target_object, *(target_objects or [])] if name]
            primary = target_object or (names[0] if names else None)
            return {
                "scope": {
                    "scope_kind": "object_set" if len(names) > 1 else "single_object",
                    "primary_target": primary,
                    "object_names": names,
                    "object_count": len(names),
                    "object_roles": [],
                },
                "pairs": [],
                "summary": {"pair_count": 0},
            }

        def get_view_diagnostics(
            self,
            target_object=None,
            target_objects=None,
            camera_name=None,
            focus_target=None,
            view_name=None,
            orbit_horizontal=0.0,
            orbit_vertical=0.0,
            zoom_factor=None,
            persist_view=False,
        ):
            names = [name for name in [target_object, *(target_objects or [])] if name]
            return {
                "view_query": {
                    "requested_view_source": "user_perspective",
                    "resolved_view_source": "user_perspective",
                    "analysis_backend": "mirrored_user_perspective",
                    "available": True,
                    "state_restored": True,
                },
                "summary": {
                    "target_count": len(names),
                    "visible_count": len(names),
                    "partially_visible_count": 0,
                    "fully_occluded_count": 0,
                    "outside_frame_count": 0,
                    "unavailable_count": 0,
                    "centered_target_count": len(names),
                    "framing_issue_count": 0,
                },
                "targets": [],
            }


    class CollectionHandler:
        def manage_collection(self, action, collection_name, new_name=None, parent_name=None, object_name=None):
            return f"Created collection '{collection_name}' under Scene Collection"


    class ModelingHandler:
        def create_primitive(self, primitive_type, radius=1.0, size=2.0, location=None, rotation=None, name=None):
            return f"Created {primitive_type} named '{name or primitive_type}'"

        def transform_object(self, name, location=None, rotation=None, scale=None):
            return f"Transformed object '{name}'"


    class MeshHandler:
        def extrude_region(self, move=None):
            return f"Extruded region by {move}"


    router_area.get_router_handler = lambda: RouterHandler()
    router_area._should_attach_repair_suggestion = lambda payload: False
    router_area._scene_has_meaningful_guided_objects = lambda: False
    scene_area.get_scene_handler = lambda: SceneHandler()
    di.get_scene_handler = lambda: SceneHandler()
    collection_area.get_collection_handler = lambda: CollectionHandler()
    modeling_area.get_modeling_handler = lambda: ModelingHandler()
    mesh_area.get_mesh_handler = lambda: MeshHandler()
    router_helper.is_router_enabled = lambda: False
    """
)


@pytest.mark.slow
def test_streamable_guided_session_expands_visible_tools_after_goal_handoff(tmp_path: Path):
    """The same streamable session should keep spatial support pinned while search can reach build tools after goal handoff."""

    script_path = write_server_script(tmp_path, _PATCHED_GUIDED_STREAMABLE_SERVER)

    async def run(url: str) -> None:
        async with streamable_client(url) as client:
            initial_tools = {tool.name for tool in await client.list_tools()}
            assert "scene_scope_graph" in initial_tools
            assert "scene_relation_graph" in initial_tools
            assert "scene_view_diagnostics" in initial_tools
            assert "modeling_create_primitive" not in initial_tools
            assert "collection_manage" not in initial_tools

            goal_result = result_payload(
                await client.call_tool(
                    "router_set_goal",
                    {"goal": "create a low-poly squirrel matching front and side reference images"},
                )
            )
            assert goal_result["guided_handoff"]["recipe_id"] == "low_poly_creature_blockout"
            assert goal_result["guided_flow_state"]["current_step"] == "bootstrap_primary_workset"
            assert goal_result["guided_flow_state"]["next_actions"] == ["create_primary_workset"]

            status_result = result_payload(await client.call_tool("router_get_status", {}))
            assert status_result["guided_flow_state"]["current_step"] == "bootstrap_primary_workset"
            assert any(
                set(rule.get("names") or ()) >= {"guided_register_part", "modeling_create_primitive"}
                for rule in status_result["visibility_rules"]
                if rule.get("components") == ["tool"] or rule.get("components") == {"tool"}
            )

            post_goal_tools = {tool.name for tool in await client.list_tools()}
            assert {"scene_scope_graph", "scene_relation_graph", "scene_view_diagnostics"}.issubset(post_goal_tools)

            blocked = result_payload(
                await client.call_tool(
                    "modeling_create_primitive",
                    {"primitive_type": "Cone", "name": "Squirrel_Ear_L", "radius": 0.1, "guided_role": "ear_pair"},
                )
            )
            assert "tool family 'secondary_parts'" in blocked

            blocked_named_body = result_payload(
                await client.call_tool(
                    "modeling_create_primitive",
                    {
                        "primitive_type": "Sphere",
                        "name": "Sphere",
                        "radius": 0.5,
                        "guided_role": "body_core",
                    },
                )
            )
            assert "Guided naming blocked object name 'Sphere'" in blocked_named_body
            assert "Body" in blocked_named_body

            await client.call_tool(
                "guided_register_part",
                {"object_name": "Squirrel_Body", "role": "body_core"},
            )
            await client.call_tool(
                "guided_register_part",
                {"object_name": "Squirrel_Head", "role": "head_mass"},
            )

            stale_status = result_payload(await client.call_tool("router_get_status", {}))
            assert stale_status["guided_flow_state"]["current_step"] == "place_secondary_parts"
            assert stale_status["guided_flow_state"]["spatial_refresh_required"] is True
            assert stale_status["guided_flow_state"]["next_actions"] == ["refresh_spatial_context"]
            assert stale_status["guided_flow_state"]["allowed_families"] == [
                "spatial_context",
                "reference_context",
            ]

            refresh_scope = {"target_object": "Squirrel_Body", "target_objects": ["Squirrel_Head"]}
            await client.call_tool("scene_scope_graph", refresh_scope)
            await client.call_tool(
                "scene_relation_graph",
                {**refresh_scope, "goal_hint": "assembled creature"},
            )
            await client.call_tool(
                "scene_view_diagnostics",
                {**refresh_scope, "view_name": "TOP"},
            )

            refreshed_status = result_payload(await client.call_tool("router_get_status", {}))
            assert refreshed_status["guided_flow_state"]["spatial_refresh_required"] is False
            assert refreshed_status["guided_flow_state"]["allowed_roles"] == [
                "tail_mass",
                "snout_mass",
                "ear_pair",
                "foreleg_pair",
                "hindleg_pair",
            ]

            weak_name_warning = result_payload(
                await client.call_tool(
                    "guided_register_part",
                    {"object_name": "ForeL", "role": "foreleg_pair"},
                )
            )
            assert weak_name_warning["guided_naming"]["status"] == "warning"
            assert "ForeLeg_L" in weak_name_warning["message"]

            ear_result = result_payload(
                await client.call_tool(
                    "modeling_create_primitive",
                    {
                        "primitive_type": "Cone",
                        "name": "Squirrel_Ear_L",
                        "location": [0.22, -0.3, 1.82],
                        "radius": 0.1,
                        "guided_role": "ear_pair",
                    },
                )
            )
            assert ear_result == "Created Cone named 'Squirrel_Ear_L'"

            primitive_search = result_payload(
                await client.call_tool(
                    "search_tools",
                    {"query": "low poly creature head body primitive blockout"},
                )
            )
            primitive_names = {item["name"] for item in primitive_search}
            assert "modeling_create_primitive" in primitive_names

            collection_search = result_payload(
                await client.call_tool(
                    "search_tools",
                    {"query": "create squirrel blockout collection move object to collection"},
                )
            )
            collection_names = {item["name"] for item in collection_search}
            assert "collection_manage" in collection_names

    with run_streamable_server(script_path) as url:
        asyncio.run(run(url))


@pytest.mark.slow
def test_streamable_guided_scene_cleanup_returns_after_goal_handoff(tmp_path: Path):
    """Build-phase cleanup should not leave Streamable HTTP clients waiting for a tool response."""

    script_path = write_server_script(tmp_path, _PATCHED_GUIDED_STREAMABLE_SERVER)

    async def run(url: str) -> None:
        async with streamable_client(url) as client:
            goal_result = result_payload(
                await client.call_tool(
                    "router_set_goal",
                    {"goal": "create a low-poly squirrel matching front and side reference images"},
                )
            )
            assert goal_result["guided_handoff"]["recipe_id"] == "low_poly_creature_blockout"

            cleanup_result = result_payload(
                await client.call_tool("scene_clean_scene", {"keep_lights_and_cameras": True})
            )
            assert cleanup_result == "Scene cleaned."

            status_result = result_payload(await client.call_tool("router_get_status", {}))
            assert status_result["guided_flow_state"]["current_step"] == "bootstrap_primary_workset"
            assert status_result["guided_flow_state"]["spatial_refresh_required"] is False
            assert status_result["guided_flow_state"]["next_actions"] == ["create_primary_workset"]

            tools_after_cleanup = {tool.name for tool in await client.list_tools()}
            assert {"scene_scope_graph", "scene_relation_graph", "scene_view_diagnostics"}.issubset(tools_after_cleanup)

    with run_streamable_server(script_path) as url:
        asyncio.run(run(url))


@pytest.mark.slow
def test_streamable_guided_dirty_mesh_tool_returns_and_rearms_spatial_context(tmp_path: Path):
    """Dirty secondary mesh tools should return and rearm spatial refresh on Streamable HTTP."""

    script_path = write_server_script(tmp_path, _PATCHED_GUIDED_STREAMABLE_SERVER)

    async def run(url: str) -> None:
        async with streamable_client(url) as client:
            await client.call_tool(
                "router_set_goal",
                {"goal": "create a low-poly squirrel matching front and side reference images"},
            )
            await client.call_tool(
                "guided_register_part",
                {"object_name": "Squirrel_Body", "role": "body_core"},
            )
            await client.call_tool(
                "guided_register_part",
                {"object_name": "Squirrel_Head", "role": "head_mass"},
            )
            refresh_scope = {"target_object": "Squirrel_Body", "target_objects": ["Squirrel_Head"]}
            await client.call_tool("scene_scope_graph", refresh_scope)
            await client.call_tool(
                "scene_relation_graph",
                {**refresh_scope, "goal_hint": "assembled creature"},
            )
            await client.call_tool(
                "scene_view_diagnostics",
                {**refresh_scope, "view_name": "TOP"},
            )

            before_status = result_payload(await client.call_tool("router_get_status", {}))
            assert before_status["guided_flow_state"]["current_step"] == "place_secondary_parts"
            assert before_status["guided_flow_state"]["spatial_refresh_required"] is False

            mesh_result = result_payload(await client.call_tool("mesh_extrude_region", {"move": [0.0, 0.0, 0.1]}))
            assert mesh_result == "Extruded region by [0.0, 0.0, 0.1]"

            after_status = result_payload(await client.call_tool("router_get_status", {}))
            assert after_status["guided_flow_state"]["current_step"] == "place_secondary_parts"
            assert after_status["guided_flow_state"]["spatial_refresh_required"] is True
            assert after_status["guided_flow_state"]["next_actions"] == ["refresh_spatial_context"]

    with run_streamable_server(script_path) as url:
        asyncio.run(run(url))


@pytest.mark.slow
def test_streamable_guided_view_diagnostics_requires_bound_scope_before_refresh_clears(tmp_path: Path):
    """A successful view-diagnostics payload should not clear the gate until scope/relation checks bind the same scope."""

    script_path = write_server_script(tmp_path, _PATCHED_GUIDED_STREAMABLE_SERVER)

    async def run(url: str) -> None:
        async with streamable_client(url) as client:
            await client.call_tool(
                "router_set_goal",
                {"goal": "create a low-poly squirrel matching front and side reference images"},
            )
            await client.call_tool(
                "guided_register_part",
                {"object_name": "Squirrel_Body", "role": "body_core"},
            )
            await client.call_tool(
                "guided_register_part",
                {"object_name": "Squirrel_Head", "role": "head_mass"},
            )

            first_view = result_payload(
                await client.call_tool(
                    "scene_view_diagnostics",
                    {"target_objects": ["Squirrel_Body", "Squirrel_Head"], "view_name": "FRONT"},
                )
            )
            assert first_view["payload"]["view_query"]["available"] is True

            stale_status = result_payload(await client.call_tool("router_get_status", {}))
            checks_by_tool = {
                item["tool_name"]: item["status"] for item in stale_status["guided_flow_state"]["required_checks"]
            }
            assert stale_status["guided_flow_state"]["spatial_refresh_required"] is True
            assert checks_by_tool["scene_view_diagnostics"] == "pending"

            refresh_scope = {"target_object": "Squirrel_Body", "target_objects": ["Squirrel_Head"]}
            await client.call_tool("scene_scope_graph", refresh_scope)
            await client.call_tool(
                "scene_relation_graph",
                {**refresh_scope, "goal_hint": "assembled creature"},
            )

            still_pending_status = result_payload(await client.call_tool("router_get_status", {}))
            checks_by_tool = {
                item["tool_name"]: item["status"]
                for item in still_pending_status["guided_flow_state"]["required_checks"]
            }
            assert checks_by_tool["scene_scope_graph"] == "completed"
            assert checks_by_tool["scene_relation_graph"] == "completed"
            assert checks_by_tool["scene_view_diagnostics"] == "pending"

            second_view = result_payload(
                await client.call_tool(
                    "scene_view_diagnostics",
                    {**refresh_scope, "view_name": "FRONT"},
                )
            )
            assert second_view["payload"]["view_query"]["available"] is True

            refreshed_status = result_payload(await client.call_tool("router_get_status", {}))
            assert refreshed_status["guided_flow_state"]["spatial_refresh_required"] is False

    with run_streamable_server(script_path) as url:
        asyncio.run(run(url))


@pytest.mark.slow
def test_streamable_guided_transform_object_fails_cleanly_during_spatial_refresh_gate(tmp_path: Path):
    """A stale direct transform attempt should fail cleanly without dropping the Streamable HTTP session."""

    script_path = write_server_script(tmp_path, _PATCHED_GUIDED_STREAMABLE_SERVER)

    async def run(url: str) -> None:
        async with streamable_client(url) as client:
            await client.call_tool(
                "router_set_goal",
                {"goal": "create a low-poly squirrel matching front and side reference images"},
            )

            body_result = result_payload(
                await client.call_tool(
                    "modeling_create_primitive",
                    {
                        "primitive_type": "Sphere",
                        "name": "Squirrel_Body",
                        "location": [0.0, 0.0, 0.6],
                        "radius": 0.5,
                        "guided_role": "body_core",
                    },
                )
            )
            assert body_result == "Created Sphere named 'Squirrel_Body'"

            head_result = result_payload(
                await client.call_tool(
                    "modeling_create_primitive",
                    {
                        "primitive_type": "Sphere",
                        "name": "Squirrel_Head",
                        "location": [0.0, -0.4, 1.45],
                        "radius": 0.35,
                        "guided_role": "head_mass",
                    },
                )
            )
            assert head_result == "Created Sphere named 'Squirrel_Head'"

            with pytest.raises(ToolError, match="Unknown tool: 'modeling_transform_object'"):
                await client.call_tool(
                    "modeling_transform_object",
                    {"name": "Squirrel_Head", "scale": [1.0, 0.92, 0.95]},
                )

            status_result = result_payload(await client.call_tool("router_get_status", {}))
            assert status_result["current_phase"] == "build"
            assert status_result["guided_flow_state"]["current_step"] == "place_secondary_parts"
            assert status_result["guided_flow_state"]["spatial_refresh_required"] is True
            assert status_result["guided_flow_state"]["next_actions"] == ["refresh_spatial_context"]

    with run_streamable_server(script_path) as url:
        asyncio.run(run(url))


@pytest.mark.slow
def test_streamable_guided_reconnect_resets_build_only_visibility_but_keeps_spatial_support(tmp_path: Path):
    """A new streamable session should keep default spatial support while build tools remain search-discovered only."""

    script_path = write_server_script(tmp_path, _PATCHED_GUIDED_STREAMABLE_SERVER)

    async def first_session(url: str) -> tuple[str | None, set[str], set[str]]:
        async with streamable_client(url) as client:
            bootstrap_tools = {tool.name for tool in await client.list_tools()}
            status_payload = result_payload(await client.call_tool("router_get_status", {}))
            goal_result = result_payload(
                await client.call_tool(
                    "router_set_goal",
                    {"goal": "create a low-poly squirrel matching front and side reference images"},
                )
            )
            assert goal_result["guided_handoff"]["recipe_id"] == "low_poly_creature_blockout"
            build_tools = {tool.name for tool in await client.list_tools()}
            return status_payload.get("session_id"), bootstrap_tools, build_tools

    async def second_session(url: str) -> tuple[str | None, set[str]]:
        async with streamable_client(url) as client:
            status_payload = result_payload(await client.call_tool("router_get_status", {}))
            bootstrap_tools = {tool.name for tool in await client.list_tools()}
            return status_payload.get("session_id"), bootstrap_tools

    with run_streamable_server(script_path) as url:
        first_session_id, first_bootstrap_tools, first_build_tools = asyncio.run(first_session(url))
        second_session_id, second_bootstrap_tools = asyncio.run(second_session(url))

    assert first_session_id is not None
    assert second_session_id is not None
    assert first_session_id != second_session_id

    assert {"scene_scope_graph", "scene_relation_graph", "scene_view_diagnostics"}.issubset(first_bootstrap_tools)
    assert {"scene_scope_graph", "scene_relation_graph", "scene_view_diagnostics"}.issubset(first_build_tools)
    assert "modeling_create_primitive" not in first_build_tools
    assert "collection_manage" not in first_build_tools

    assert {"scene_scope_graph", "scene_relation_graph", "scene_view_diagnostics"}.issubset(second_bootstrap_tools)
    assert "modeling_create_primitive" not in second_bootstrap_tools
    assert "collection_manage" not in second_bootstrap_tools
