import asyncio
import os
import sys

# Set UTF-8 encoding for stdout
if sys.platform == "win32":
    import io
    sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding='utf-8')

from autogen_ext.models.openai import AzureOpenAIChatCompletionClient
from autogen_ext.tools.mcp import McpWorkbench, StdioServerParams
from autogen_agentchat.agents import AssistantAgent

model_client = AzureOpenAIChatCompletionClient(
    model="gpt-4.1-nano",
    api_version="2024-12-01-preview",
    azure_endpoint="https://agent-gallery108.openai.azure.com/",
    api_key="FfMrurWuud7oU0LFBc2hbtHCJsrcASg3DhQWs0i8qk34HZjQP9ojJQQJ99CAACfhMk5XJ3w3AAAAACOGxnqu",
)

async def run_test(agent, task_name, task):
    print(f"\n{'='*50}")
    print(f"TEST: {task_name}")
    print(f"{'='*50}")
    
    result = await agent.run(task=task)
    final_message = result.messages[-1]
    print(final_message.content)
    print(f"\n{task_name} - COMPLETED\n")

async def main():
    # Use the current venv's Python to run the CLI directly
    import sys as sys_module
    
    python_exe = sys_module.executable
    
    server_params = StdioServerParams(
        command="pipx",
        args=["run", "jmeter-mcp-server"]
    )

    async with McpWorkbench(server_params) as workbench:
        agent = AssistantAgent(
            name="assistant",
            model_client=model_client,
            workbench=workbench,
            system_message="You are a JMeter performance testing expert. Use MCP tools to analyze and correlate test files.",
        )
            
        result = await agent.run(task="What are the tools available?")
        final_message = result.messages[-1]
        print(final_message.content)

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