import asyncio

from dotenv import load_dotenv
from mcp import ClientSession
from mcp.client.streamable_http import streamable_http_client
from pydantic_ai import Agent, Tool

from ez_ptc import Toolkit

load_dotenv()

MCP_URL = "http://localhost:8001/mcp"


async def main():
    import json
    async with streamable_http_client(MCP_URL) as (read, write, _):
        async with ClientSession(read, write) as session:
            await session.initialize()

            toolkit = await Toolkit.from_mcp(
                session,
                assist_tool_chaining=False,
            )
            print(f"Discovered {len(toolkit)} tools from CardDekho MCP server\n")
            print(toolkit.tool_prompt())

            execute_fn = toolkit.as_tool()
            pydantic_tool = Tool(execute_fn, takes_ctx=False)

            agent = Agent(
                "openai:gpt-4.1",
                system_prompt=(
                    "Assist the user with car queries using the provided tools.\n\n"
                    # + toolkit.tool_prompt()
                ),
                tools=[pydantic_tool],
            )

            print("CardDekho Car Assistant (type 'quit' to exit)")
            print("-" * 45)
            while True:
                query = input("\nYou: ").strip()
                if not query:
                    continue
                if query.lower() in ("quit", "exit", "q"):
                    print("Goodbye!")
                    break
                result = await agent.run(query)
                print(json.dumps(json.loads(result.all_messages_json()), indent=4))
                print(f"\nAssistant: {result.output}")


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