"""RAG agent: Qdrant vector store + Ollama LLM + LangGraph."""
import os
from langchain_ollama import ChatOllama, OllamaEmbeddings
from langchain_qdrant import QdrantVectorStore
from langchain_core.documents import Document
from langchain.tools import tool
from langgraph.prebuilt import create_react_agent

QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost")
QDRANT_PORT = int(os.getenv("QDRANT_PORT", "6333"))
COLLECTION_NAME = "knowledge"

embeddings = OllamaEmbeddings(model="nomic-embed-text")
_vs = None


def get_vector_store() -> QdrantVectorStore:
    """Return singleton Qdrant vector store, creating collection if needed."""
    global _vs
    if _vs is None:
        from qdrant_client import QdrantClient
        from qdrant_client.models import Distance, VectorParams
        client = QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT)
        existing = [c.name for c in client.get_collections().collections]
        if COLLECTION_NAME not in existing:
            client.create_collection(
                collection_name=COLLECTION_NAME,
                vectors_config=VectorParams(size=768, distance=Distance.COSINE),
            )
        _vs = QdrantVectorStore(
            client=client,
            collection_name=COLLECTION_NAME,
            embedding=embeddings,
        )
    return _vs


@tool
def search_knowledge_base(query: str) -> str:
    """Search the knowledge base for relevant passages.

    Args:
        query: search query string
    """
    docs = get_vector_store().similarity_search(query, k=5)
    if not docs:
        return "No results found."
    return "\n\n".join(f"{i+1}. {d.page_content}" for i, d in enumerate(docs))


@tool
def add_to_knowledge_base(text: str) -> str:
    """Add a text passage to the knowledge base.

    Args:
        text: text to store
    """
    get_vector_store().add_documents([Document(page_content=text)])
    return "Added to knowledge base."

llm = ChatOllama(model="llama3", temperature=0.0)

agent = create_react_agent(
    model=llm,
    tools=[search_knowledge_base, add_to_knowledge_base],
    state_modifier=(
        "You are a helpful assistant with access to a local knowledge base. "
        "Use search_knowledge_base to find information. "
        "Use add_to_knowledge_base to store new information."
    ),
)


def run_agent(user_input: str) -> str:
    """Run the agent and return the last message content."""
    from langchain_core.messages import HumanMessage
    result = agent.invoke({"messages": [HumanMessage(content=user_input)]})
    return result["messages"][-1].content