import asyncio
import json
import sys

from mcp import ClientSession, StdioServerParameters, stdio_client


PLUGIN_DIR = r"C:\Users\darsh.shah\.codex\plugins\cache\openai-curated\gpt-researcher\f78e3ad49297672a905eb7afb6aa0cef34edc79e"

try:
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
except AttributeError:
    pass


async def main() -> None:
    params = StdioServerParameters(
        command="cmd",
        args=["/c", PLUGIN_DIR + r"\run-gptr-mcp.cmd"],
    )

    async with stdio_client(params) as (read_stream, write_stream):
        async with ClientSession(read_stream, write_stream) as session:
            init_result = await session.initialize()
            print("initialized")
            print(init_result.serverInfo.name)

            tools = await session.list_tools()
            tool_names = [tool.name for tool in tools.tools]
            print(json.dumps(tool_names))

            mode = sys.argv[1] if len(sys.argv) > 1 else "quick"
            query = (
                sys.argv[2]
                if len(sys.argv) > 2
                else "persistent LLM memory systems risks"
            )

            if mode == "quick":
                quick = await session.call_tool("quick_search", {"query": query})
                print("quick_search_ok")
                print(str(quick)[:2000])
                return

            if mode == "deep":
                deep = await session.call_tool("deep_research", {"query": query})
                print("deep_research_ok")
                print(str(deep)[:4000])
                follow_up = await session.call_tool(
                    "quick_search", {"query": "memory poisoning in LLM agents"}
                )
                print("follow_up_quick_search_ok")
                print(str(follow_up)[:1000])
                return

            raise SystemExit(f"Unknown mode: {mode}")


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