# server.py

from mcp.server.fastmcp import FastMCP
import loguru
import time
import json

mcp = FastMCP("test-server", "0.0.1")


@mcp.tool()
def print_data(data: str) -> str:
    """
    Print the data.
    """
    loguru.logger.warning(f"Received data: {data} (type: {type(data)})")
    return f"data: {data}"


if __name__ == "__main__":
    import sys
    transport = sys.argv[1] if len(sys.argv) > 1 else "stdio"
    if transport != "stdio" and transport != "streamable-http":
        raise ValueError("Transport must be 'stdio' or 'streamable-http'")
    mcp.run(transport=transport)

# stdio_client.py
from mcp import ClientSession, StdioServerParameters, types
from mcp.client.stdio import stdio_client
import json

server_params = StdioServerParameters(
    command="python",
    args=["-u", "server.py"],
    env=None,
)


async def run():
    async with stdio_client(server_params) as (read, write):
        async with ClientSession(read, write) as session:
            await session.initialize()
            args = {"data": json.dumps({"message": "hello"})}
            result = await session.call_tool("print_data", arguments=args)
            print("Tool call result:", result)


if __name__ == "__main__":
    import asyncio
    asyncio.run(run())
    
# http_client.py
from mcp.client.streamable_http import streamablehttp_client
from mcp import ClientSession
import json

async def main():
    async with streamablehttp_client("http://localhost:8000/mcp") as (
        read_stream,
        write_stream,
        _,
    ):
        async with ClientSession(read_stream, write_stream) as session:
            await session.initialize()
            args = {"data": json.dumps({"message": "hello"})}
            tool_result = await session.call_tool("print_data", arguments=args)
            print("Tool call result:", tool_result)

if __name__ == "__main__":
    import asyncio
    asyncio.run(main())