"""Interactive LangChain agent with local RAG memory on Qdrant and Ollama."""

import os
from functools import lru_cache
from typing import Any
from uuid import uuid4

from langchain.agents import create_agent
from langchain_core.documents import Document
from langchain_core.tools import tool
from langchain_ollama import ChatOllama, OllamaEmbeddings
from langchain_qdrant import QdrantVectorStore
from langchain_text_splitters import RecursiveCharacterTextSplitter
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams

QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333")
QDRANT_COLLECTION = os.getenv("QDRANT_COLLECTION", "rag_memory")
OLLAMA_CHAT_MODEL = os.getenv("OLLAMA_CHAT_MODEL", "llama3")
OLLAMA_EMBED_MODEL = os.getenv("OLLAMA_EMBED_MODEL", "nomic-embed-text")
EMBEDDING_SIZE = int(os.getenv("OLLAMA_EMBEDDING_SIZE", "768"))


@lru_cache(maxsize=1)
def get_embeddings() -> OllamaEmbeddings:
    return OllamaEmbeddings(model=OLLAMA_EMBED_MODEL)


@lru_cache(maxsize=1)
def get_vector_store() -> QdrantVectorStore:
    client = QdrantClient(url=QDRANT_URL)
    collections = {item.name for item in client.get_collections().collections}
    if QDRANT_COLLECTION not in collections:
        client.create_collection(
            collection_name=QDRANT_COLLECTION,
            vectors_config=VectorParams(size=EMBEDDING_SIZE, distance=Distance.COSINE),
        )
    return QdrantVectorStore(
        client=client,
        collection_name=QDRANT_COLLECTION,
        embedding=get_embeddings(),
    )


def chunk_document(content: str, title: str) -> list[Document]:
    splitter = RecursiveCharacterTextSplitter(chunk_size=700, chunk_overlap=100)
    return splitter.create_documents([content], metadatas=[{"title": title}])


@tool
def search_knowledge_base(query: str, max_results: int = 3) -> str:
    """Search relevant chunks in the local Qdrant knowledge base."""
    results = get_vector_store().similarity_search_with_score(query, k=max_results)
    if not results:
        return "No relevant documents found."

    lines: list[str] = []
    for index, (document, score) in enumerate(results, start=1):
        title = document.metadata.get("title", "untitled")
        snippet = document.page_content.replace("\n", " ")[:300]
        lines.append(f"{index}. {title} (score={score:.4f}): {snippet}")
    return "\n".join(lines)


@tool
def add_to_knowledge_base(content: str, title: str) -> str:
    """Split content into chunks and store it in the local Qdrant knowledge base."""
    documents = chunk_document(content, title)
    ids = [str(uuid4()) for _ in documents]
    get_vector_store().add_documents(documents, ids=ids)
    return f"Added {len(documents)} chunk(s) from '{title}' to the knowledge base."


def build_agent() -> Any:
    llm = ChatOllama(model=OLLAMA_CHAT_MODEL)
    return create_agent(
        model=llm,
        tools=[search_knowledge_base, add_to_knowledge_base],
        system_prompt=(
            "You are a RAG assistant. Use search_knowledge_base before answering "
            "questions that may depend on stored knowledge, and use "
            "add_to_knowledge_base when the user asks to remember information."
        ),
    )


def run_cli() -> None:
    print("RAG agent. Commands: /add <title> | <content>, /search <query>, /quit")
    agent = None

    while True:
        try:
            user_input = input(">>> ").strip()
        except EOFError:
            break

        if not user_input:
            continue
        if user_input == "/quit":
            break
        if user_input.startswith("/add "):
            payload = user_input[5:]
            if "|" in payload:
                title, content = [part.strip() for part in payload.split("|", 1)]
            else:
                title, content = "note", payload.strip()
            print(add_to_knowledge_base.invoke({"content": content, "title": title}))
            continue
        if user_input.startswith("/search "):
            query = user_input[8:].strip()
            print(search_knowledge_base.invoke({"query": query, "max_results": 3}))
            continue

        if agent is None:
            agent = build_agent()
        response = agent.invoke({"messages": [{"role": "user", "content": user_input}]})
        messages = response.get("messages", [])
        print(messages[-1].content if messages else response)


if __name__ == "__main__":
    run_cli()