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():
    server_params = StdioServerParams(
        command="pipx",
        args=["run", "jmeter-mcp"]
    )

    async with McpWorkbench(server_params) as workbench:
        agent = AssistantAgent(
            name="assistant",
            model_client=model_client,
            workbench=workbench,
            system_message="Use tools to solve tasks.",
        )

        # Test 1: Compare and find dynamic values
        print("\n" + "="*50)
        print("TEST: Compare Test1.xml and Test2.xml for Dynamic Values")
        print("="*50)
        
        # Read the test files
        test1_path = r"C:\Users\jjeswant\Downloads\Test1.xml"
        test2_path = r"C:\Users\jjeswant\Downloads\Test2.xml"
        
        try:
            with open(test1_path, 'r', encoding='utf-8', errors='ignore') as f:
                test1_content = f.read()
            with open(test2_path, 'r', encoding='utf-8', errors='ignore') as f:
                test2_content = f.read()
            
            print(f"\n[✓] Loaded Test1.xml ({len(test1_content)} bytes)")
            print(f"[✓] Loaded Test2.xml ({len(test2_content)} bytes)")
            
            # Create task to analyze the files
            task = f"""
Compare these two JMeter XML test result files and identify dynamic values that need correlation:

Test1.xml (first 1000 chars):
{test1_content[:1000]}...

Test2.xml (first 1000 chars):
{test2_content[:1000]}...
"""
            
            await run_test(agent, "Find Dynamic Values Correlation", task)
            
        except FileNotFoundError as e:
            print(f"\n[ERROR] Error reading test files: {e}")
            print(f"   Expected files:")
            print(f"   - {test1_path}")
            print(f"   - {test2_path}")
        except Exception as e:
            print(f"\n[ERROR] Error during comparison: {e}")


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