"""Tests for response caching middleware."""

import sys
import tempfile
import warnings
from pathlib import Path
from unittest.mock import AsyncMock, MagicMock

import mcp.types
import pytest
from inline_snapshot import snapshot
from key_value.aio.stores.filetree import (
    FileTreeStore,
    FileTreeV1CollectionSanitizationStrategy,
    FileTreeV1KeySanitizationStrategy,
)
from key_value.aio.stores.memory import MemoryStore
from key_value.aio.wrappers.statistics.wrapper import (
    GetStatistics,
    KVStoreCollectionStatistics,
    PutStatistics,
)
from mcp.server.lowlevel.helper_types import ReadResourceContents
from mcp.types import TextContent, TextResourceContents
from pydantic import AnyUrl, BaseModel

from fastmcp import Context, FastMCP
from fastmcp.client.client import CallToolResult, Client
from fastmcp.client.transports import FastMCPTransport
from fastmcp.prompts.function_prompt import FunctionPrompt
from fastmcp.prompts.prompt import Message, Prompt
from fastmcp.resources.resource import Resource
from fastmcp.server.middleware.caching import (
    CachableToolResult,
    CallToolSettings,
    ResponseCachingMiddleware,
    ResponseCachingStatistics,
)
from fastmcp.server.middleware.middleware import CallNext, MiddlewareContext
from fastmcp.tools.tool import Tool, ToolResult

TEST_URI = AnyUrl("https://test_uri")

SAMPLE_READ_RESOURCE_CONTENTS = ReadResourceContents(
    content="test_text",
    mime_type="text/plain",
)


def sample_resource_fn() -> list[ReadResourceContents]:
    return [SAMPLE_READ_RESOURCE_CONTENTS]


def sample_prompt_fn() -> Message:
    return Message("test_text")


SAMPLE_RESOURCE = Resource.from_function(
    fn=sample_resource_fn, uri=TEST_URI, name="test_resource"
)

SAMPLE_PROMPT = Prompt.from_function(fn=sample_prompt_fn, name="test_prompt")
SAMPLE_GET_PROMPT_RESULT = mcp.types.GetPromptResult(
    messages=[Message("test_text").to_mcp_prompt_message()]
)
SAMPLE_TOOL = Tool(name="test_tool", parameters={"param1": "value1", "param2": 42})
SAMPLE_TOOL_RESULT = ToolResult(
    content=[TextContent(type="text", text="test_text")],
    structured_content={"result": "test_result"},
)
SAMPLE_TOOL_RESULT_LARGE = ToolResult(
    content=[TextContent(type="text", text="test_text" * 100)],
    structured_content={"result": "test_result"},
)


class CrazyModel(BaseModel):
    a: int
    b: int
    c: str
    d: float
    e: bool
    f: list[int]
    g: dict[str, int]
    h: list[dict[str, int]]
    i: dict[str, list[int]]


@pytest.fixture
def crazy_model() -> CrazyModel:
    return CrazyModel(
        a=5,
        b=10,
        c="test",
        d=1.0,
        e=True,
        f=[1, 2, 3],
        g={"a": 1, "b": 2},
        h=[{"a": 1, "b": 2}],
        i={"a": [1, 2]},
    )


class TrackingCalculator:
    add_calls: int
    multiply_calls: int
    crazy_calls: int
    very_large_response_calls: int

    def __init__(self):
        self.add_calls = 0
        self.multiply_calls = 0
        self.crazy_calls = 0
        self.very_large_response_calls = 0

    def add(self, a: int, b: int) -> int:
        self.add_calls += 1
        return a + b

    def multiply(self, a: int, b: int) -> int:
        self.multiply_calls += 1
        return a * b

    def very_large_response(self) -> str:
        self.very_large_response_calls += 1
        return "istenchars" * 100000  # 1,000,000 characters, 1mb

    def crazy(self, a: CrazyModel) -> CrazyModel:
        self.crazy_calls += 1
        return a

    def how_to_calculate(self, a: int, b: int) -> str:
        return f"To calculate {a} + {b}, you need to add {a} and {b} together."

    def get_add_calls(self) -> str:
        return str(self.add_calls)

    def get_multiply_calls(self) -> str:
        return str(self.multiply_calls)

    def get_crazy_calls(self) -> str:
        return str(self.crazy_calls)

    async def update_tool_list(self, context: Context):
        import mcp.types

        await context.send_notification(mcp.types.ToolListChangedNotification())

    def add_tools(self, fastmcp: FastMCP, prefix: str = ""):
        _ = fastmcp.add_tool(tool=Tool.from_function(fn=self.add, name=f"{prefix}add"))
        _ = fastmcp.add_tool(
            tool=Tool.from_function(fn=self.multiply, name=f"{prefix}multiply")
        )
        _ = fastmcp.add_tool(
            tool=Tool.from_function(fn=self.crazy, name=f"{prefix}crazy")
        )
        _ = fastmcp.add_tool(
            tool=Tool.from_function(
                fn=self.very_large_response, name=f"{prefix}very_large_response"
            )
        )
        _ = fastmcp.add_tool(
            tool=Tool.from_function(
                fn=self.update_tool_list, name=f"{prefix}update_tool_list"
            )
        )

    def add_prompts(self, fastmcp: FastMCP, prefix: str = ""):
        _ = fastmcp.add_prompt(
            prompt=FunctionPrompt.from_function(
                fn=self.how_to_calculate, name=f"{prefix}how_to_calculate"
            )
        )

    def add_resources(self, fastmcp: FastMCP, prefix: str = ""):
        _ = fastmcp.add_resource(
            resource=Resource.from_function(
                fn=self.get_add_calls,
                uri="resource://add_calls",
                name=f"{prefix}add_calls",
            )
        )
        _ = fastmcp.add_resource(
            resource=Resource.from_function(
                fn=self.get_multiply_calls,
                uri="resource://multiply_calls",
                name=f"{prefix}multiply_calls",
            )
        )
        _ = fastmcp.add_resource(
            resource=Resource.from_function(
                fn=self.get_crazy_calls,
                uri="resource://crazy_calls",
                name=f"{prefix}crazy_calls",
            )
        )


@pytest.fixture
def tracking_calculator() -> TrackingCalculator:
    return TrackingCalculator()


@pytest.fixture
def mock_context() -> MiddlewareContext[mcp.types.CallToolRequestParams]:
    """Create a mock middleware context for tool calls."""
    context = MagicMock(spec=MiddlewareContext[mcp.types.CallToolRequestParams])
    context.message = mcp.types.CallToolRequestParams(
        name="test_tool", arguments={"param1": "value1", "param2": 42}
    )
    context.method = "tools/call"
    return context


@pytest.fixture
def mock_call_next() -> CallNext[mcp.types.CallToolRequestParams, ToolResult]:
    """Create a mock call_next function."""
    return AsyncMock(
        return_value=ToolResult(
            content=[TextContent(type="text", text="test result")],
            structured_content={"result": "success", "value": 123},
        )
    )


@pytest.fixture
def sample_tool_result() -> ToolResult:
    """Create a sample tool result for testing."""
    return ToolResult(
        content=[TextContent(type="text", text="cached result")],
        structured_content={"cached": True, "data": "test"},
    )


class TestResponseCachingMiddleware:
    """Test ResponseCachingMiddleware functionality."""

    def test_initialization(self):
        """Test middleware initialization."""
        assert ResponseCachingMiddleware(
            call_tool_settings=CallToolSettings(
                included_tools=["tool1"],
                excluded_tools=["tool2"],
            ),
        )

    @pytest.mark.parametrize(
        ("tool_name", "included_tools", "excluded_tools", "result"),
        [
            ("tool", ["tool", "tool2"], [], True),
            ("tool", ["second tool", "third tool"], [], False),
            ("tool", [], ["tool"], False),
            ("tool", [], ["second tool"], True),
            ("tool", ["tool", "second tool"], ["tool"], False),
            ("tool", ["tool", "second tool"], ["second tool"], True),
        ],
        ids=[
            "tool is included",
            "tool is not included",
            "tool is excluded",
            "tool is not excluded",
            "tool is included and excluded (excluded takes precedence)",
            "tool is included and not excluded",
        ],
    )
    def test_tool_call_filtering(
        self,
        tool_name: str,
        included_tools: list[str],
        excluded_tools: list[str],
        result: bool,
    ):
        """Test tool filtering logic."""

        middleware1 = ResponseCachingMiddleware(
            call_tool_settings=CallToolSettings(
                included_tools=included_tools, excluded_tools=excluded_tools
            ),
        )
        assert middleware1._matches_tool_cache_settings(tool_name=tool_name) is result


@pytest.mark.skipif(
    sys.platform == "win32",
    reason="SQLite caching tests are flaky on Windows due to temp directory issues.",
)
class TestResponseCachingMiddlewareIntegration:
    """Integration tests with real FastMCP server."""

    @pytest.fixture(params=["memory", "filetree"])
    async def caching_server(
        self,
        tracking_calculator: TrackingCalculator,
        request: pytest.FixtureRequest,
    ):
        """Create a FastMCP server for caching tests."""
        mcp = FastMCP("CachingTestServer", dereference_schemas=False)

        with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as temp_dir:
            with warnings.catch_warnings():
                warnings.simplefilter("ignore", UserWarning)
                file_store = FileTreeStore(
                    data_directory=Path(temp_dir),
                    key_sanitization_strategy=FileTreeV1KeySanitizationStrategy(
                        Path(temp_dir)
                    ),
                    collection_sanitization_strategy=FileTreeV1CollectionSanitizationStrategy(
                        Path(temp_dir)
                    ),
                )
            response_caching_middleware = ResponseCachingMiddleware(
                cache_storage=file_store
                if request.param == "filetree"
                else MemoryStore(),
            )

            mcp.add_middleware(middleware=response_caching_middleware)

            tracking_calculator.add_tools(fastmcp=mcp)
            tracking_calculator.add_resources(fastmcp=mcp)
            tracking_calculator.add_prompts(fastmcp=mcp)

            yield mcp

    @pytest.fixture
    def non_caching_server(self, tracking_calculator: TrackingCalculator):
        """Create a FastMCP server for non-caching tests."""
        mcp = FastMCP("NonCachingTestServer")
        tracking_calculator.add_tools(fastmcp=mcp)
        return mcp

    async def test_list_tools(
        self, caching_server: FastMCP, tracking_calculator: TrackingCalculator
    ):
        """Test that tool list caching works with a real FastMCP server."""

        async with Client(caching_server) as client:
            pre_tool_list: list[mcp.types.Tool] = await client.list_tools()
            assert len(pre_tool_list) == 5

            # Add a tool and make sure it's missing from the list tool response
            _ = caching_server.add_tool(
                tool=Tool.from_function(fn=tracking_calculator.add, name="add_2")
            )

            post_tool_list: list[mcp.types.Tool] = await client.list_tools()
            assert len(post_tool_list) == 5

            assert pre_tool_list == post_tool_list

    async def test_call_tool(
        self,
        caching_server: FastMCP,
        tracking_calculator: TrackingCalculator,
    ):
        """Test that caching works with a real FastMCP server."""
        tracking_calculator.add_tools(fastmcp=caching_server)

        async with Client[FastMCPTransport](transport=caching_server) as client:
            call_tool_result_one: CallToolResult = await client.call_tool(
                "add", {"a": 5, "b": 3}
            )

            assert tracking_calculator.add_calls == 1
            call_tool_result_two: CallToolResult = await client.call_tool(
                "add", {"a": 5, "b": 3}
            )
            assert call_tool_result_one == call_tool_result_two

    async def test_call_tool_very_large_value(
        self,
        caching_server: FastMCP,
        tracking_calculator: TrackingCalculator,
    ):
        """Test that caching works with a real FastMCP server."""
        tracking_calculator.add_tools(fastmcp=caching_server)

        async with Client[FastMCPTransport](transport=caching_server) as client:
            call_tool_result_one: CallToolResult = await client.call_tool(
                "very_large_response", {}
            )

            assert tracking_calculator.very_large_response_calls == 1
            call_tool_result_two: CallToolResult = await client.call_tool(
                "very_large_response", {}
            )
            assert call_tool_result_one == call_tool_result_two
            assert tracking_calculator.very_large_response_calls == 2

    async def test_call_tool_crazy_value(
        self,
        caching_server: FastMCP,
        tracking_calculator: TrackingCalculator,
        crazy_model: CrazyModel,
    ):
        """Test that caching works with a real FastMCP server."""
        tracking_calculator.add_tools(fastmcp=caching_server)

        async with Client[FastMCPTransport](transport=caching_server) as client:
            call_tool_result_one: CallToolResult = await client.call_tool(
                "crazy", {"a": crazy_model}
            )

            assert tracking_calculator.crazy_calls == 1
            call_tool_result_two: CallToolResult = await client.call_tool(
                "crazy", {"a": crazy_model}
            )
            assert call_tool_result_one == call_tool_result_two
            assert tracking_calculator.crazy_calls == 1

    async def test_list_resources(
        self, caching_server: FastMCP, tracking_calculator: TrackingCalculator
    ):
        """Test that list resources caching works with a real FastMCP server."""
        async with Client[FastMCPTransport](transport=caching_server) as client:
            pre_resource_list: list[mcp.types.Resource] = await client.list_resources()

            assert len(pre_resource_list) == 3

            tracking_calculator.add_resources(fastmcp=caching_server)

            post_resource_list: list[mcp.types.Resource] = await client.list_resources()
            assert len(post_resource_list) == 3

            assert pre_resource_list == post_resource_list

    async def test_read_resource(
        self, caching_server: FastMCP, tracking_calculator: TrackingCalculator
    ):
        """Test that get resources caching works with a real FastMCP server."""
        async with Client[FastMCPTransport](transport=caching_server) as client:
            pre_resource = await client.read_resource(uri="resource://add_calls")
            assert isinstance(pre_resource[0], TextResourceContents)
            assert pre_resource[0].text == "0"

            tracking_calculator.add_calls = 1

            post_resource = await client.read_resource(uri="resource://add_calls")
            assert isinstance(post_resource[0], TextResourceContents)
            assert post_resource[0].text == "0"
            assert pre_resource == post_resource

    async def test_list_prompts(
        self, caching_server: FastMCP, tracking_calculator: TrackingCalculator
    ):
        """Test that list prompts caching works with a real FastMCP server."""
        async with Client[FastMCPTransport](transport=caching_server) as client:
            pre_prompt_list: list[mcp.types.Prompt] = await client.list_prompts()

            assert len(pre_prompt_list) == 1

            tracking_calculator.add_prompts(fastmcp=caching_server)

            post_prompt_list: list[mcp.types.Prompt] = await client.list_prompts()

            assert len(post_prompt_list) == 1

            assert pre_prompt_list == post_prompt_list

    async def test_get_prompts(
        self, caching_server: FastMCP, tracking_calculator: TrackingCalculator
    ):
        """Test that get prompts caching works with a real FastMCP server."""
        async with Client[FastMCPTransport](transport=caching_server) as client:
            pre_prompt = await client.get_prompt(
                name="how_to_calculate", arguments={"a": 5, "b": 3}
            )

            pre_prompt_content = pre_prompt.messages[0].content
            assert isinstance(pre_prompt_content, TextContent)
            assert (
                pre_prompt_content.text
                == "To calculate 5 + 3, you need to add 5 and 3 together."
            )

            tracking_calculator.add_prompts(fastmcp=caching_server)

            post_prompt = await client.get_prompt(
                name="how_to_calculate", arguments={"a": 5, "b": 3}
            )

            assert pre_prompt == post_prompt

    async def test_statistics(
        self,
        caching_server: FastMCP,
    ):
        """Test that statistics are collected correctly."""
        caching_middleware = caching_server.middleware[0]
        assert isinstance(caching_middleware, ResponseCachingMiddleware)

        async with Client[FastMCPTransport](transport=caching_server) as client:
            statistics = caching_middleware.statistics()
            assert statistics == snapshot(ResponseCachingStatistics())

            _ = await client.call_tool("add", {"a": 5, "b": 3})

            statistics = caching_middleware.statistics()
            assert statistics == snapshot(
                ResponseCachingStatistics(
                    list_tools=KVStoreCollectionStatistics(
                        get=GetStatistics(count=2, hit=1, miss=1),
                        put=PutStatistics(count=1),
                    ),
                    call_tool=KVStoreCollectionStatistics(
                        get=GetStatistics(count=1, miss=1), put=PutStatistics(count=1)
                    ),
                )
            )

            _ = await client.call_tool("add", {"a": 5, "b": 3})

            statistics = caching_middleware.statistics()
            assert statistics == snapshot(
                ResponseCachingStatistics(
                    list_tools=KVStoreCollectionStatistics(
                        get=GetStatistics(count=2, hit=1, miss=1),
                        put=PutStatistics(count=1),
                    ),
                    call_tool=KVStoreCollectionStatistics(
                        get=GetStatistics(count=2, hit=1, miss=1),
                        put=PutStatistics(count=1),
                    ),
                )
            )


class TestCachableToolResult:
    def test_wrap_and_unwrap(self):
        tool_result = ToolResult(
            "unstructured content",
            structured_content={"structured": "content"},
            meta={"meta": "data"},
        )

        cached_tool_result = CachableToolResult.wrap(tool_result).unwrap()

        assert cached_tool_result.content == tool_result.content
        assert cached_tool_result.structured_content == tool_result.structured_content
        assert cached_tool_result.meta == tool_result.meta


class TestCachingWithImportedServerPrefixes:
    """Test that caching preserves prefixes from imported servers.

    Regression tests for issue #2300: ResponseCachingMiddleware was losing
    prefix information when caching components from imported servers.
    """

    @pytest.fixture
    async def parent_with_imported_child(self, tracking_calculator: TrackingCalculator):
        """Create a parent server with an imported child server (prefixed)."""
        child = FastMCP("child")
        tracking_calculator.add_tools(fastmcp=child)
        tracking_calculator.add_resources(fastmcp=child)
        tracking_calculator.add_prompts(fastmcp=child)

        parent = FastMCP("parent")
        parent.add_middleware(ResponseCachingMiddleware())
        parent.mount(child, namespace="child")

        return parent

    async def test_tool_prefixes_preserved_after_cache_hit(
        self, parent_with_imported_child: FastMCP
    ):
        """Tool names should retain prefix after being served from cache."""
        async with Client(parent_with_imported_child) as client:
            # First call populates cache
            tools_first = await client.list_tools()
            tool_names_first = [t.name for t in tools_first]

            # Second call should come from cache
            tools_cached = await client.list_tools()
            tool_names_cached = [t.name for t in tools_cached]

            # All tools should have prefix in both calls
            assert all(name.startswith("child_") for name in tool_names_first)
            assert all(name.startswith("child_") for name in tool_names_cached)
            assert tool_names_first == tool_names_cached

    async def test_resource_prefixes_preserved_after_cache_hit(
        self, parent_with_imported_child: FastMCP
    ):
        """Resource URIs should retain prefix after being served from cache."""
        async with Client(parent_with_imported_child) as client:
            # First call populates cache
            resources_first = await client.list_resources()
            resource_uris_first = [str(r.uri) for r in resources_first]

            # Second call should come from cache
            resources_cached = await client.list_resources()
            resource_uris_cached = [str(r.uri) for r in resources_cached]

            # All resources should have prefix in URI path in both calls
            # Resources get path-style prefix: resource://child/path
            assert all("://child/" in uri for uri in resource_uris_first)
            assert all("://child/" in uri for uri in resource_uris_cached)
            assert resource_uris_first == resource_uris_cached

    async def test_prompt_prefixes_preserved_after_cache_hit(
        self, parent_with_imported_child: FastMCP
    ):
        """Prompt names should retain prefix after being served from cache."""
        async with Client(parent_with_imported_child) as client:
            # First call populates cache
            prompts_first = await client.list_prompts()
            prompt_names_first = [p.name for p in prompts_first]

            # Second call should come from cache
            prompts_cached = await client.list_prompts()
            prompt_names_cached = [p.name for p in prompts_cached]

            # All prompts should have prefix in both calls
            assert all(name.startswith("child_") for name in prompt_names_first)
            assert all(name.startswith("child_") for name in prompt_names_cached)
            assert prompt_names_first == prompt_names_cached

    async def test_prefixed_tool_callable_after_cache_hit(
        self,
        parent_with_imported_child: FastMCP,
        tracking_calculator: TrackingCalculator,
    ):
        """Prefixed tools should be callable after cache populates."""
        async with Client(parent_with_imported_child) as client:
            # Trigger cache population
            await client.list_tools()
            await client.list_tools()  # From cache

            # Tool should be callable with prefixed name
            result = await client.call_tool("child_add", {"a": 5, "b": 3})
            assert not result.is_error
            assert tracking_calculator.add_calls == 1
