from langchain_community.vectorstores import FAISS
from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings


class KeywordEmbeddings(Embeddings):
    vocabulary = ["faiss", "index", "retrieval", "metadata", "save", "reuse"]

    def _embed(self, text: str) -> list[float]:
        lowered = text.lower()
        return [float(lowered.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)


documents = [
    Document(
        page_content="Build the FAISS index from embedded LangChain documents.",
        metadata={"source": "indexing"},
    ),
    Document(
        page_content="Use the retriever view when a RAG chain needs documents.",
        metadata={"source": "retrieval"},
    ),
    Document(
        page_content="Save the FAISS index and docstore before reuse.",
        metadata={"source": "persistence"},
    ),
    Document(
        page_content="Attach metadata filters before narrowing search results.",
        metadata={"source": "metadata"},
    ),
]

embeddings = KeywordEmbeddings()
vector_store = FAISS.from_documents(documents, embeddings)
print(f"created index documents: {vector_store.index.ntotal}")

match = vector_store.similarity_search("FAISS index", k=1)[0]
print(f"best match: {match.page_content}")

vector_store.save_local("faiss_index")
loaded_store = FAISS.load_local(
    "faiss_index",
    embeddings,
    allow_dangerous_deserialization=True,
)
loaded_match = loaded_store.similarity_search("save index for reuse", k=1)[0]
print(f"loaded match: {loaded_match.page_content}")

retriever = loaded_store.as_retriever(search_kwargs={"k": 1})
retrieved = retriever.invoke("retrieval documents")[0]
print(f"retriever match: {retrieved.page_content}")