from langchain_community.vectorstores import Chroma
from langchain_ollama import OllamaEmbeddings
from langchain.schema import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter

COLLECTION_NAME = "rag_collection"
EMBEDDING_MODEL = "nomic-embed-text"
CHROMA_PERSIST_DIR = "./chroma_db"


class ChromaStore:
    def __init__(self, collection_name: str = COLLECTION_NAME):
        self.client = Chroma(
            embedding_function=OllamaEmbeddings(model=EMBEDDING_MODEL),
            collection_name=collection_name,
            persist_directory=CHROMA_PERSIST_DIR,
        )

    def add_documents(self, docs: list[Document]) -> None:
        self.client.add_documents(docs)

    def search(self, query: str, limit: int = 5) -> list[Document]:
        return self.client.similarity_search(query, k=limit)


def add_documents(content: str, title: str) -> int:
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,
        chunk_overlap=50,
        separators=["\n\n", "\n", ".", " ", ""],
    )
    chunks = splitter.split_text(content)
    docs = [
        Document(
            page_content=chunk,
            metadata={"title": title, "chunk_index": i, "source": title},
        )
        for i, chunk in enumerate(chunks)
    ]
    store = ChromaStore()
    store.add_documents(docs)
    return len(docs)


def search_documents(query: str, max_results: int = 5) -> list[dict]:
    store = ChromaStore()
    results = store.search(query, limit=max_results)
    output = []
    for doc in results:
        output.append(
            {
                "content": doc.page_content,
                "metadata": doc.metadata,
            }
        )
    return output