from pathlib import Path
from shutil import rmtree

from langchain_chroma import Chroma
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings


class KeywordEmbeddings(Embeddings):
    vocabulary = (
        "chroma",
        "persistent",
        "collection",
        "password",
        "invoice",
        "billing",
    )

    def _embed(self, text: str) -> list[float]:
        lower_text = text.lower()
        return [float(lower_text.count(term)) for term in self.vocabulary]

    def embed_documents(self, texts: list[str]) -> list[list[float]]:
        return [self._embed(text) for text in texts]

    def embed_query(self, text: str) -> list[float]:
        return self._embed(text)


persist_directory = Path("chroma_langchain_db")
if persist_directory.exists():
    rmtree(persist_directory)

collection_name = "support_knowledge"
embeddings = KeywordEmbeddings()
documents = [
    Document(
        page_content=(
            "Chroma stores embedded support notes in a local persistent collection."
        ),
        metadata={"source": "architecture"},
    ),
    Document(
        page_content="Password reset tickets should route to the identity queue.",
        metadata={"source": "accounts"},
    ),
    Document(
        page_content="Invoice export failures should route to billing operations.",
        metadata={"source": "billing"},
    ),
]

vector_store = Chroma(
    collection_name=collection_name,
    embedding_function=embeddings,
    persist_directory=str(persist_directory),
)
vector_store.add_documents(
    documents=documents,
    ids=["architecture-chroma", "accounts-password", "billing-invoice"],
)

stored_ids = vector_store.get()["ids"]
reopened_store = Chroma(
    collection_name=collection_name,
    embedding_function=embeddings,
    persist_directory=str(persist_directory),
)
reopened_ids = reopened_store.get()["ids"]

query = "persistent Chroma collection"
top_match = reopened_store.similarity_search(query, k=1)[0]

print(f"collection: {collection_name}")
print(f"documents stored: {len(stored_ids)}")
print(f"reopened documents: {len(reopened_ids)}")
print(f"query: {query}")
print(f"top match source: {top_match.metadata['source']}")
print(f"top match text: {top_match.page_content}")