import json
from contextlib import asynccontextmanager

import pytest
from mcp.server.lowlevel.helper_types import ReadResourceContents
from mcp.types import TextContent, TextResourceContents

from fastmcp import FastMCP
from fastmcp.client import Client
from fastmcp.client.transports import FastMCPTransport
from fastmcp.exceptions import NotFoundError
from fastmcp.server.proxy import FastMCPProxy


class TestBasicMount:
    """Test basic mounting functionality."""

    async def test_mount_simple_server(self):
        """Test mounting a simple server and accessing its tool."""
        # Create main app and sub-app
        main_app = FastMCP("MainApp")
        sub_app = FastMCP("SubApp")

        # Add a tool to the sub-app
        @sub_app.tool()
        def sub_tool() -> str:
            return "This is from the sub app"

        # Mount the sub-app to the main app
        main_app.mount("sub", sub_app)

        # Get tools from main app, should include sub_app's tools
        tools = await main_app.get_tools()
        assert "sub_sub_tool" in tools

        async with Client(main_app) as client:
            result = await client.call_tool("sub_sub_tool", {})
            assert isinstance(result[0], TextContent)
            assert result[0].text == "This is from the sub app"

    async def test_mount_with_custom_separator(self):
        """Test mounting with a custom tool separator (deprecated but still supported)."""
        main_app = FastMCP("MainApp")
        sub_app = FastMCP("SubApp")

        @sub_app.tool()
        def greet(name: str) -> str:
            return f"Hello, {name}!"

        # Mount without custom separator - custom separators are deprecated
        main_app.mount("sub", sub_app)

        # Tool should be accessible with the default separator
        tools = await main_app.get_tools()
        assert "sub_greet" in tools

        # Call the tool
        result = await main_app._mcp_call_tool("sub_greet", {"name": "World"})
        assert isinstance(result[0], TextContent)
        assert result[0].text == "Hello, World!"

    async def test_mount_invalid_resource_prefix(self):
        main_app = FastMCP("MainApp")
        api_app = FastMCP("APIApp")

        # This test doesn't apply anymore with the new prefix format
        # just mount the server to maintain test coverage
        main_app.mount("api:sub", api_app)

    async def test_mount_invalid_resource_separator(self):
        main_app = FastMCP("MainApp")
        api_app = FastMCP("APIApp")

        # This test doesn't apply anymore with the new prefix format
        # Mount without deprecated parameters
        main_app.mount("api", api_app)

    async def test_unmount_server(self):
        """Test unmounting a server removes access to its tools."""
        main_app = FastMCP("MainApp")
        sub_app = FastMCP("SubApp")

        @sub_app.tool()
        def sub_tool() -> str:
            return "This is from the sub app"

        # Mount the sub-app
        main_app.mount("sub", sub_app)

        # Verify it was mounted
        tools = await main_app.get_tools()
        assert "sub_sub_tool" in tools

        # Unmount the sub-app
        main_app.unmount("sub")

        # Verify it was unmounted
        tools = await main_app.get_tools()
        assert "sub_sub_tool" not in tools

        # Calling the tool should fail
        with pytest.raises(NotFoundError, match="Unknown tool: sub_sub_tool"):
            await main_app._mcp_call_tool("sub_sub_tool", {})

    async def test_mount_with_no_prefix(self):
        main_app = FastMCP("MainApp")
        sub_app = FastMCP("SubApp")

        @sub_app.tool()
        def sub_tool() -> str:
            return "This is from the sub app"

        # Mount with empty prefix but without deprecated separators
        main_app.mount(prefix="", server=sub_app)

        tools = await main_app.get_tools()
        # With empty prefix, the format is now "_sub_tool" instead of "sub_tool"
        assert "_sub_tool" in tools


class TestMultipleServerMount:
    """Test mounting multiple servers simultaneously."""

    async def test_mount_multiple_servers(self):
        """Test mounting multiple servers with different prefixes."""
        main_app = FastMCP("MainApp")
        weather_app = FastMCP("WeatherApp")
        news_app = FastMCP("NewsApp")

        @weather_app.tool()
        def get_forecast() -> str:
            return "Weather forecast"

        @news_app.tool()
        def get_headlines() -> str:
            return "News headlines"

        # Mount both apps
        main_app.mount("weather", weather_app)
        main_app.mount("news", news_app)

        # Check both are accessible
        tools = await main_app.get_tools()
        assert "weather_get_forecast" in tools
        assert "news_get_headlines" in tools

        # Call tools from both mounted servers
        result1 = await main_app._mcp_call_tool("weather_get_forecast", {})
        assert isinstance(result1[0], TextContent)
        assert result1[0].text == "Weather forecast"

        result2 = await main_app._mcp_call_tool("news_get_headlines", {})
        assert isinstance(result2[0], TextContent)
        assert result2[0].text == "News headlines"

    async def test_mount_same_prefix(self):
        """Test that mounting with the same prefix replaces the previous mount."""
        main_app = FastMCP("MainApp")
        first_app = FastMCP("FirstApp")
        second_app = FastMCP("SecondApp")

        @first_app.tool()
        def first_tool() -> str:
            return "First app tool"

        @second_app.tool()
        def second_tool() -> str:
            return "Second app tool"

        # Mount first app
        main_app.mount("api", first_app)
        tools = await main_app.get_tools()
        assert "api_first_tool" in tools

        # Mount second app with same prefix
        main_app.mount("api", second_app)
        tools = await main_app.get_tools()

        # First app's tool should no longer be accessible
        assert "api_first_tool" not in tools

        # Second app's tool should be accessible
        assert "api_second_tool" in tools


class TestDynamicChanges:
    """Test that changes to mounted servers are reflected dynamically."""

    async def test_adding_tool_after_mounting(self):
        """Test that tools added after mounting are accessible."""
        main_app = FastMCP("MainApp")
        sub_app = FastMCP("SubApp")

        # Mount the sub-app before adding any tools
        main_app.mount("sub", sub_app)

        # Initially, there should be no tools from sub_app
        tools = await main_app.get_tools()
        assert not any(key.startswith("sub_") for key in tools)

        # Add a tool to the sub-app after mounting
        @sub_app.tool()
        def dynamic_tool() -> str:
            return "Added after mounting"

        # The tool should be accessible through the main app
        tools = await main_app.get_tools()
        assert "sub_dynamic_tool" in tools

        # Call the dynamically added tool
        result = await main_app._mcp_call_tool("sub_dynamic_tool", {})
        assert isinstance(result[0], TextContent)
        assert result[0].text == "Added after mounting"

    async def test_removing_tool_after_mounting(self):
        """Test that tools removed from mounted servers are no longer accessible."""
        main_app = FastMCP("MainApp")
        sub_app = FastMCP("SubApp")

        @sub_app.tool()
        def temp_tool() -> str:
            return "Temporary tool"

        # Mount the sub-app
        main_app.mount("sub", sub_app)

        # Initially, the tool should be accessible
        tools = await main_app.get_tools()
        assert "sub_temp_tool" in tools

        # Remove the tool from sub_app
        sub_app._tool_manager._tools.pop("temp_tool")

        # The tool should no longer be accessible
        # Refresh the cache by clearing it
        main_app._cache.cache.clear()
        tools = await main_app.get_tools()
        assert "sub_temp_tool" not in tools


class TestResourcesAndTemplates:
    """Test mounting with resources and resource templates."""

    async def test_mount_with_resources(self):
        """Test mounting a server with resources."""
        main_app = FastMCP("MainApp")
        data_app = FastMCP("DataApp")

        @data_app.resource(uri="data://users")
        async def get_users():
            return ["user1", "user2"]

        # Mount the data app
        main_app.mount("data", data_app)

        # Resource should be accessible through main app
        resources = await main_app.get_resources()
        assert "data://data/users" in resources

        # Check that resource can be accessed
        async with Client(main_app) as client:
            result = await client.read_resource("data://data/users")
            assert isinstance(result[0], TextResourceContents)
            assert json.loads(result[0].text) == ["user1", "user2"]

    async def test_mount_with_resource_templates(self):
        """Test mounting a server with resource templates."""
        main_app = FastMCP("MainApp")
        user_app = FastMCP("UserApp")

        @user_app.resource(uri="users://{user_id}/profile")
        def get_user_profile(user_id: str) -> dict:
            return {"id": user_id, "name": f"User {user_id}"}

        # Mount the user app
        main_app.mount("api", user_app)

        # Template should be accessible through main app
        templates = await main_app.get_resource_templates()
        assert "users://api/{user_id}/profile" in templates

        # Check template instantiation
        async with Client(main_app) as client:
            result = await client.read_resource("users://api/123/profile")
            assert isinstance(result[0], TextResourceContents)
            profile = json.loads(result[0].text)
            assert profile["id"] == "123"
            assert profile["name"] == "User 123"

    async def test_adding_resource_after_mounting(self):
        """Test adding a resource after mounting."""
        main_app = FastMCP("MainApp")
        data_app = FastMCP("DataApp")

        # Mount the data app before adding resources
        main_app.mount("data", data_app)

        # Add a resource after mounting
        @data_app.resource(uri="data://config")
        def get_config():
            return {"version": "1.0"}

        # Resource should be accessible through main app
        resources = await main_app.get_resources()
        assert "data://data/config" in resources

        # Check access to the resource
        async with Client(main_app) as client:
            result = await client.read_resource("data://data/config")
            assert isinstance(result[0], TextResourceContents)
            config = json.loads(result[0].text)
            assert config["version"] == "1.0"


class TestPrompts:
    """Test mounting with prompts."""

    async def test_mount_with_prompts(self):
        """Test mounting a server with prompts."""
        main_app = FastMCP("MainApp")
        assistant_app = FastMCP("AssistantApp")

        @assistant_app.prompt()
        def greeting(name: str) -> str:
            return f"Hello, {name}!"

        # Mount the assistant app
        main_app.mount("assistant", assistant_app)

        # Prompt should be accessible through main app
        prompts = await main_app.get_prompts()
        assert "assistant_greeting" in prompts

        # Render the prompt
        result = await main_app._mcp_get_prompt("assistant_greeting", {"name": "World"})
        assert result.messages is not None
        # The message should contain our greeting text

    async def test_adding_prompt_after_mounting(self):
        """Test adding a prompt after mounting."""
        main_app = FastMCP("MainApp")
        assistant_app = FastMCP("AssistantApp")

        # Mount the assistant app before adding prompts
        main_app.mount("assistant", assistant_app)

        # Add a prompt after mounting
        @assistant_app.prompt()
        def farewell(name: str) -> str:
            return f"Goodbye, {name}!"

        # Prompt should be accessible through main app
        prompts = await main_app.get_prompts()
        assert "assistant_farewell" in prompts

        # Render the prompt
        result = await main_app._mcp_get_prompt("assistant_farewell", {"name": "World"})
        assert result.messages is not None
        # The message should contain our farewell text


class TestProxyServer:
    """Test mounting a proxy server."""

    async def test_mount_proxy_server(self):
        """Test mounting a proxy server."""
        # Create original server
        original_server = FastMCP("OriginalServer")

        @original_server.tool()
        def get_data(query: str) -> str:
            return f"Data for {query}"

        # Create proxy server
        proxy_server = FastMCP.as_proxy(
            Client(transport=FastMCPTransport(original_server))
        )

        # Mount proxy server
        main_app = FastMCP("MainApp")
        main_app.mount("proxy", proxy_server)

        # Tool should be accessible through main app
        tools = await main_app.get_tools()
        assert "proxy_get_data" in tools

        # Call the tool
        result = await main_app._mcp_call_tool("proxy_get_data", {"query": "test"})
        assert isinstance(result[0], TextContent)
        assert result[0].text == "Data for test"

    async def test_dynamically_adding_to_proxied_server(self):
        """Test that changes to the original server are reflected in the mounted proxy."""
        # Create original server
        original_server = FastMCP("OriginalServer")

        # Create proxy server
        proxy_server = FastMCP.as_proxy(
            Client(transport=FastMCPTransport(original_server))
        )

        # Mount proxy server
        main_app = FastMCP("MainApp")
        main_app.mount("proxy", proxy_server)

        # Add a tool to the original server
        @original_server.tool()
        def dynamic_data() -> str:
            return "Dynamic data"

        # Tool should be accessible through main app via proxy
        tools = await main_app.get_tools()
        assert "proxy_dynamic_data" in tools

        # Call the tool
        result = await main_app._mcp_call_tool("proxy_dynamic_data", {})
        assert isinstance(result[0], TextContent)
        assert result[0].text == "Dynamic data"

    async def test_proxy_server_with_resources(self):
        """Test mounting a proxy server with resources."""
        # Create original server
        original_server = FastMCP("OriginalServer")

        @original_server.resource(uri="config://settings")
        def get_config():
            return {"api_key": "12345"}

        # Create proxy server
        proxy_server = FastMCP.as_proxy(
            Client(transport=FastMCPTransport(original_server))
        )

        # Mount proxy server
        main_app = FastMCP("MainApp")
        main_app.mount("proxy", proxy_server)

        # Resource should be accessible through main app
        result = await main_app._mcp_read_resource("config://proxy/settings")
        assert isinstance(result[0], ReadResourceContents)
        config = json.loads(result[0].content)
        assert config["api_key"] == "12345"

    async def test_proxy_server_with_prompts(self):
        """Test mounting a proxy server with prompts."""
        # Create original server
        original_server = FastMCP("OriginalServer")

        @original_server.prompt()
        def welcome(name: str) -> str:
            return f"Welcome, {name}!"

        # Create proxy server
        proxy_server = FastMCP.as_proxy(
            Client(transport=FastMCPTransport(original_server))
        )

        # Mount proxy server
        main_app = FastMCP("MainApp")
        main_app.mount("proxy", proxy_server)

        # Prompt should be accessible through main app
        result = await main_app._mcp_get_prompt("proxy_welcome", {"name": "World"})
        assert result.messages is not None
        # The message should contain our welcome text


class TestAsProxyKwarg:
    """Test the as_proxy kwarg."""

    async def test_as_proxy_defaults_false(self):
        mcp = FastMCP("Main")
        sub = FastMCP("Sub")

        mcp.mount("sub", sub)

        assert mcp._mounted_servers["sub"].server is sub

    async def test_as_proxy_false(self):
        mcp = FastMCP("Main")
        sub = FastMCP("Sub")

        mcp.mount("sub", sub, as_proxy=False)

        assert mcp._mounted_servers["sub"].server is sub

    async def test_as_proxy_true(self):
        mcp = FastMCP("Main")
        sub = FastMCP("Sub")

        mcp.mount("sub", sub, as_proxy=True)

        assert mcp._mounted_servers["sub"].server is not sub
        assert isinstance(mcp._mounted_servers["sub"].server, FastMCPProxy)

    async def test_as_proxy_defaults_true_if_lifespan(self):
        @asynccontextmanager
        async def lifespan(mcp: FastMCP):
            yield

        mcp = FastMCP("Main")
        sub = FastMCP("Sub", lifespan=lifespan)

        mcp.mount("sub", sub)

        assert mcp._mounted_servers["sub"].server is not sub
        assert isinstance(mcp._mounted_servers["sub"].server, FastMCPProxy)

    async def test_as_proxy_ignored_for_proxy_mounts_default(self):
        mcp = FastMCP("Main")
        sub = FastMCP("Sub")
        sub_proxy = FastMCP.as_proxy(Client(transport=FastMCPTransport(sub)))

        mcp.mount("sub", sub_proxy)

        assert mcp._mounted_servers["sub"].server is sub_proxy

    async def test_as_proxy_ignored_for_proxy_mounts_false(self):
        mcp = FastMCP("Main")
        sub = FastMCP("Sub")
        sub_proxy = FastMCP.as_proxy(Client(transport=FastMCPTransport(sub)))

        mcp.mount("sub", sub_proxy, as_proxy=False)

        assert mcp._mounted_servers["sub"].server is sub_proxy

    async def test_as_proxy_ignored_for_proxy_mounts_true(self):
        mcp = FastMCP("Main")
        sub = FastMCP("Sub")
        sub_proxy = FastMCP.as_proxy(Client(transport=FastMCPTransport(sub)))

        mcp.mount("sub", sub_proxy, as_proxy=True)

        assert mcp._mounted_servers["sub"].server is sub_proxy

    async def test_as_proxy_mounts_still_have_live_link(self):
        mcp = FastMCP("Main")
        sub = FastMCP("Sub")

        mcp.mount("sub", sub, as_proxy=True)

        assert len(await mcp.get_tools()) == 0

        @sub.tool()
        def hello():
            return "hi"

        assert len(await mcp.get_tools()) == 1

    async def test_sub_lifespan_is_executed(self):
        lifespan_check = []

        @asynccontextmanager
        async def lifespan(mcp: FastMCP):
            lifespan_check.append("start")
            yield

        mcp = FastMCP("Main")
        sub = FastMCP("Sub", lifespan=lifespan)

        @sub.tool()
        def hello():
            return "hi"

        mcp.mount("sub", sub, as_proxy=True)

        assert lifespan_check == []

        async with Client(mcp) as client:
            await client.call_tool("sub_hello", {})

        assert lifespan_check == ["start"]