import asyncio
import random
import json
import os
from datetime import timedelta
from mcp.client.stdio import stdio_client
from mcp import ClientSession, StdioServerParameters
import dtlpy as dl

dl.setenv('prod')
if dl.token_expired():
    dl.login()

# Create server parameters for stdio connection
server_params = StdioServerParameters(
    command="e:\\Applications\\dtlpy-mcp\\.venv\\Scripts\\python.exe",  # Executable
    args=["-m", "dtlpymcp", "start", "-s", "tests/assets/sources.json"],  # Command line arguments
    env={"DATALOOP_API_KEY": dl.token()},  # Optional environment variables
    cwd=os.getcwd()
)

async def test_health_check():
    print("[TEST CLIENT] Connecting to MCP server and calling test tool...")
    async with stdio_client(server=server_params) as (read, write):
        async with ClientSession(read, write, read_timeout_seconds=timedelta(seconds=60)) as session:
            await session.initialize()
            tools = await session.list_tools()
            print([tool.name for tool in tools.tools])
            tool_result = await session.call_tool("dl-internal-chat-prod__list_repos")
            print("[RESULT]", tool_result)            
            tool_result = await session.call_tool("dl-internal-chat-prod__ping")
            print("[RESULT]", tool_result)

if __name__ == "__main__":
    asyncio.run(test_health_check()) 
    # asyncio.run(test_ask_dataloop()) 