import asyncio
import os
from dotenv import load_dotenv
import openai
from agents import Agent, Runner, AsyncOpenAI, OpenAIChatCompletionsModel
from agents.mcp import MCPServerSse
import json

# Load environment
load_dotenv()

# Set API key
openai.api_key = os.getenv("NEBIUS_API_KEY")
openai.base_url = "https://api.studio.nebius.com/v1/"

# MCP tool URL (replace with actual URL you got from step 4.3)
TOOL_URL = "http://127.0.0.1:7860/gradio_api/mcp/sse"

async def main():
    mcp_server = MCPServerSse({"url": TOOL_URL})

    try:
        await mcp_server.connect()

        # Create an agent and register the connected tool
        # ✅ Replace "MCP Agent" with a descriptive name if desired
        # ✅ Update instructions to guide the agent's behavior
        agent = Agent(
            name="MCP Agent",
            instructions="You help the user complete tasks using the connected tool.",
            mcp_servers=[mcp_server],
            model=OpenAIChatCompletionsModel( 
                model="Qwen/Qwen3-235B-A22B",
                openai_client=AsyncOpenAI(base_url="https://api.studio.nebius.com/v1/", api_key=os.getenv("NEBIUS_API_KEY"))
            )
        )

        # 🔁 Replace this task with your actual prompt or goal
        task = "You are a helpful assistant. Use the connected tool to assist with tasks. Read and analyse the tool output, focusing on key information and help the user in the following task:" \
        "I want to travel from Hyderabad to Phuket from 5th July to 23rd July in 2025, please help me with my flight search, what are my cheapest and fastest options."

        result = await Runner.run(agent, task)

        print(result.final_output)
        # print(json.dumps(result.new_items, indent=4))
        # print(result.raw_responses)

        # for r in result:
        #     print(r)
        #     print("\n============\n")

    finally:
        await mcp_server.cleanup()

# Run the async main function
if __name__ == "__main__":
    asyncio.run(main())