#!/usr/bin/env python3
import asyncio
import json
from typing import Any

from mcp.types import CallToolResult, TextContent, Tool as MCPTool
from agents.run_context import RunContextWrapper
from agents.mcp import MCPServer, MCPUtil


class TestMCPServer(MCPServer):
    def __init__(self, use_structured_content: bool = False):
        super().__init__(use_structured_content=use_structured_content)
        self._server_name = "test_server"
    
    async def cleanup(self) -> None:
        pass
    
    async def connect(self) -> None:
        pass
    
    async def get_prompt(self, name: str, arguments: dict[str, Any] | None = None):
        raise NotImplementedError()
    
    async def list_prompts(self, run_context: RunContextWrapper[Any], agent):
        return []
    
    @property
    def name(self) -> str:
        return self._server_name
    
    async def list_tools(self, run_context: RunContextWrapper[Any], agent) -> list:
        return [MCPTool(name="search_users", description="test", inputSchema={"type": "object"})]
    
    async def call_tool(self, tool_name: str, arguments: dict[str, Any] | None) -> CallToolResult:
        return CallToolResult(
            content=[TextContent(text="Found 2 users", type="text")],
            structuredContent={
                "users": [{"id": 1, "name": "Alice"}, {"id": 2, "name": "Bob"}],
                "total": 2
            }
        )


async def test_structured_content_fix():
    print("Testing MCP Structured Content Fix (Issue #1236)")
    
    run_context = RunContextWrapper(context=None)
    tool = MCPTool(name="search_users", description="test", inputSchema={"type": "object"})
    
    # Test 1: use_structured_content=False (returns text content)
    print("\nTest 1: use_structured_content=False")
    server_text = TestMCPServer(use_structured_content=False)
    result = await MCPUtil.invoke_mcp_tool(server_text, tool, run_context, "{}")
    parsed = json.loads(result)
    print(f"Result: {result}")
    print(f"PASS: Returns text content" if "text" in parsed else "FAIL: Expected text content")
    
    # Test 2: use_structured_content=True (THE FIX - returns structured content)
    print("\nTest 2: use_structured_content=True (THE FIX)")
    print("Before fix: Would return text content (unreachable path)")
    print("After fix: Returns structured content exclusively")
    
    server_structured = TestMCPServer(use_structured_content=True)
    result_structured = await MCPUtil.invoke_mcp_tool(server_structured, tool, run_context, "{}")
    parsed_structured = json.loads(result_structured)
    print(f"Result: {result_structured}")
    
    if "users" in parsed_structured and "text" not in parsed_structured:
        print(f"PASS: Returns structured content exclusively")
        print(f"Found {len(parsed_structured['users'])} users")
        print("FIX CONFIRMED: Previously unreachable code now works!")
    else:
        print("FAIL: Expected structured content or found text mixing")
    
    # Test 3: Fallback when no structured content
    print("\nTest 3: Fallback when no structured content")
    
    class FallbackServer(TestMCPServer):
        async def call_tool(self, tool_name: str, arguments: dict[str, Any] | None):
            return CallToolResult(
                content=[TextContent(text="No structured data", type="text")],
                structuredContent=None
            )
    
    server_fallback = FallbackServer(use_structured_content=True)
    result_fallback = await MCPUtil.invoke_mcp_tool(server_fallback, tool, run_context, "{}")
    parsed_fallback = json.loads(result_fallback)
    print(f"Result: {result_fallback}")
    print("PASS: Fallback to text content works" if parsed_fallback.get("text") == "No structured data" else "FAIL: Fallback behavior not working")


if __name__ == "__main__":
    asyncio.run(test_structured_content_fix())