import os
import httpx
from fastmcp import FastMCP

LIGHTRAG_URL = os.getenv("LIGHTRAG_URL", "http://localhost:9621")

mcp = FastMCP("LightRAG")


@mcp.tool
async def insert_documents(documents: list[str]) -> str:
    """Insert text documents into LightRAG for indexing.

    Args:
        documents: List of document strings to index. Each string is treated as a separate document.

    Returns:
        Tracking ID for the insertion operation.
    """
    async with httpx.AsyncClient(timeout=120.0) as client:
        r = await client.post(
            f"{LIGHTRAG_URL}/documents/texts",
            json={"texts": documents},
        )
        r.raise_for_status()
        data = r.json()
        return data.get("track_id", data.get("message", "unknown"))


@mcp.tool
async def query_documents(query: str, mode: str = "mix", top_k: int = 60) -> dict:
    """Query LightRAG and retrieve relevant context without LLM generation.

    Args:
        query: The search query string.
        mode: Retrieval mode - "local", "global", "hybrid", "naive", "mix" (default: "mix").
        top_k: Number of top results to retrieve (default: 60).

    Returns:
        Structured retrieval data including entities, relationships, and text chunks.
    """
    async with httpx.AsyncClient(timeout=120.0) as client:
        r = await client.post(
            f"{LIGHTRAG_URL}/query/data",
            json={
                "query": query,
                "mode": mode,
                "only_need_context": True,
                "top_k": top_k,
            },
        )
        r.raise_for_status()
        return r.json()


if __name__ == "__main__":
    mcp.run()