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_text_splitters import RecursiveCharacterTextSplitter
from langchain_qdrant import QdrantVectorStore
from qdrant_client import QdrantClient

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()
    # Ensure collection exists
    client.recreate_collection(
        collection_name="rag_memory",
        vectors_config={"size": EMBEDDING_SIZE, "distance": "cosine"},
    )
    return QdrantVectorStore(client=client, collection_name="rag_memory")


def chunk_document(content: str, title: str) -> list[Document]:
    splitter = RecursiveCharacterTextSplitter(chunk_size=500, 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."""
    vector_store = get_vector_store()
    results = vector_store.similarity_search_with_score(query, k=max_results)
    if not results:
        return "No relevant documents found."

    lines: list[str] = []
    for idx, (doc, score) in enumerate(results, start=1):
        title = doc.metadata.get("title", "untitled") if doc.metadata else "untitled"
        snippet = doc.page_content.replace("\n", " ")[:300]
        lines.append(f"{idx}. {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]
    embeddings = get_embeddings().embed_documents([doc.page_content for doc in documents])
    vector_store = get_vector_store()
    vector_store.add_documents(documents, ids=ids, embeddings=embeddings)
    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()