import asyncio
import os
from mcp import ClientSession, StdioServerParameters
from mcp.client.stdio import stdio_client


async def main():
    server = StdioServerParameters(
        command="uv",
        args=[
            "run",
            "--directory",
            "/home/df/projects/rags/lightrag-mcp",
            "python",
            "main.py",
        ],
    )

    async with stdio_client(server) as (read, write):
        async with ClientSession(read, write) as session:
            await session.initialize()

            print("--- INSERT ---")
            result = await session.call_tool(
                "insert_documents",
                arguments={
                    "documents": [
                        "Python is a high-level programming language known for its simplicity and readability.",
                        "JavaScript was created in 1995 by Brendan Eich at Netscape.",
                        "Machine learning is a subset of artificial intelligence that enables systems to learn from data.",
                    ]
                },
            )
            print(f"Insert result: {result.content[0].text[:200]}")

            print("\n--- QUERY (mix) ---")
            result = await session.call_tool(
                "query_documents",
                arguments={
                    "query": "Tell me about programming languages",
                    "mode": "mix",
                    "top_k": 60,
                },
            )
            import json

            data = json.loads(result.content[0].text)
            d = data.get("data", {})
            print(f"Entities: {len(d.get('entities', []))}")
            print(f"Relationships: {len(d.get('relationships', []))}")
            print(f"Chunks: {len(d.get('chunks', []))}")
            for c in d.get("chunks", [])[:2]:
                print(f"  - {c.get('content', '')[:100]}")

            print("\n--- QUERY (local) ---")
            result = await session.call_tool(
                "query_documents",
                arguments={"query": "What is Python?", "mode": "local", "top_k": 60},
            )
            data = json.loads(result.content[0].text)
            d = data.get("data", {})
            print(f"Entities: {len(d.get('entities', []))}")
            print(f"Chunks: {len(d.get('chunks', []))}")

            print("\n--- QUERY (global) ---")
            result = await session.call_tool(
                "query_documents",
                arguments={
                    "query": "What topics are covered?",
                    "mode": "global",
                    "top_k": 60,
                },
            )
            data = json.loads(result.content[0].text)
            d = data.get("data", {})
            print(f"Entities: {len(d.get('entities', []))}")
            print(f"Relationships: {len(d.get('relationships', []))}")

    print("\nDone!")


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