from pathlib import Path

from langchain_core.documents import Document
from langchain_core.embeddings import Embeddings
from langchain_core.vectorstores import InMemoryVectorStore
from langchain_text_splitters import RecursiveCharacterTextSplitter


class ReviewEmbeddings(Embeddings):
    vocabulary = ("refund", "billing", "approval", "customer", "ignore")

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


source_files = ["refund-policy.txt", "untrusted-note.txt"]

documents = [
    Document(
        page_content=Path(file_name).read_text(encoding="utf-8"),
        metadata={"source": file_name, "stage": "loaded"},
    )
    for file_name in source_files
]

splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=0)
chunks = splitter.split_documents(documents)

for number, chunk in enumerate(chunks, start=1):
    chunk.metadata["chunk_id"] = f"chunk-{number}"
    chunk.metadata["stage"] = "split"
    chunk.metadata["review_status"] = "needs_review"

review_decisions = {
    "refund-policy.txt": ("approved", "finance approval rule verified"),
    "untrusted-note.txt": ("rejected", "prompt-like instruction"),
}

reviewed_chunks = []
for chunk in chunks:
    decision, reason = review_decisions[chunk.metadata["source"]]
    chunk.metadata["review_status"] = decision
    chunk.metadata["review_reason"] = reason
    chunk.metadata["reviewer"] = "content-review"
    reviewed_chunks.append(chunk)

approved_chunks = [
    chunk for chunk in reviewed_chunks if chunk.metadata["review_status"] == "approved"
]
rejected_chunks = [
    chunk for chunk in reviewed_chunks if chunk.metadata["review_status"] == "rejected"
]

vector_store = InMemoryVectorStore.from_documents(
    approved_chunks,
    embedding=ReviewEmbeddings(),
)

match = vector_store.similarity_search("refund approval for billing credit", k=1)[0]

print(f"loaded documents: {len(documents)}")
print(f"split chunks: {len(chunks)}")
print(f"review queue: {len(reviewed_chunks)}")
print(
    "approved for indexing: "
    + ", ".join(chunk.metadata["source"] for chunk in approved_chunks)
)
print(
    "rejected before indexing: "
    + ", ".join(chunk.metadata["source"] for chunk in rejected_chunks)
)
print(f"indexed chunks: {len(approved_chunks)}")
print(f"search result source: {match.metadata['source']}")
print(f"search result review_status: {match.metadata['review_status']}")