import httpx
import asyncio


async def main():
    base_url = "http://localhost:9621"

    async with httpx.AsyncClient(timeout=120.0) as client:
        print("--- INSERT ---")
        docs = [
            "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.",
            "LightRAG combines knowledge graph and vector retrieval for enhanced RAG applications.",
            "FastMCP is a framework for building MCP servers in Python.",
        ]
        r = await client.post(f"{base_url}/documents/texts", json={"texts": docs})
        r.raise_for_status()
        print(f"Inserted: {r.json()}")

        print("\n--- QUERY (mix mode) ---")
        r = await client.post(
            f"{base_url}/query/data",
            json={
                "query": "Tell me about programming languages",
                "mode": "mix",
                "only_need_context": True,
                "top_k": 60,
            },
        )
        r.raise_for_status()
        result = r.json()
        print(f"mode=mix keys: {list(result.keys())}")
        if "chunks" in result:
            print(f"  chunks: {len(result['chunks'])} returned")
            for c in result["chunks"][:2]:
                print(f"    - {c.get('content', '')[:100]}...")

        print("\n--- QUERY (local mode) ---")
        r = await client.post(
            f"{base_url}/query/data",
            json={
                "query": "What is Python?",
                "mode": "local",
                "only_need_context": True,
                "top_k": 60,
            },
        )
        r.raise_for_status()
        result = r.json()
        print(f"mode=local keys: {list(result.keys())}")
        if "chunks" in result:
            print(f"  chunks: {len(result['chunks'])} returned")

        print("\n--- QUERY (global mode) ---")
        r = await client.post(
            f"{base_url}/query/data",
            json={
                "query": "What topics are covered?",
                "mode": "global",
                "only_need_context": True,
                "top_k": 60,
            },
        )
        r.raise_for_status()
        result = r.json()
        print(f"mode=global keys: {list(result.keys())}")
        if "entities" in result:
            print(f"  entities: {len(result['entities'])} returned")
        if "relationships" in result:
            print(f"  relationships: {len(result['relationships'])} returned")

    print("\nDone!")


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