from haystack import Document, Pipeline
from haystack.components.joiners import DocumentJoiner
from haystack.components.retrievers.in_memory import (
    InMemoryBM25Retriever,
    InMemoryEmbeddingRetriever,
)
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack_integrations.components.embedders.sentence_transformers import (
    SentenceTransformersDocumentEmbedder,
    SentenceTransformersTextEmbedder,
)


documents = [
    Document(
        id="hybrid-search",
        content="Hybrid retrieval combines keyword matching with semantic search.",
    ),
    Document(
        id="bm25-keywords",
        content="BM25 finds exact terms such as product names and error codes.",
    ),
    Document(
        id="vector-meaning",
        content="Embedding search finds related meaning despite different wording.",
    ),
    Document(
        id="coffee-notes",
        content="Coffee brewing depends on water temperature and grind size.",
    ),
]

model = "sentence-transformers/all-MiniLM-L6-v2"
document_embedder = SentenceTransformersDocumentEmbedder(model=model)
document_embedder.warm_up()
embedded_documents = document_embedder.run(documents=documents)["documents"]

document_store = InMemoryDocumentStore(
    embedding_similarity_function="cosine"
)
document_store.write_documents(embedded_documents)

text_embedder = SentenceTransformersTextEmbedder(model=model)
bm25_retriever = InMemoryBM25Retriever(
    document_store=document_store,
    top_k=2,
)
embedding_retriever = InMemoryEmbeddingRetriever(
    document_store=document_store,
    top_k=2,
)
joiner = DocumentJoiner(
    join_mode="reciprocal_rank_fusion",
    top_k=3,
)

pipeline = Pipeline()
pipeline.add_component("text_embedder", text_embedder)
pipeline.add_component("bm25_retriever", bm25_retriever)
pipeline.add_component("embedding_retriever", embedding_retriever)
pipeline.add_component("joiner", joiner)
pipeline.connect(
    "text_embedder.embedding",
    "embedding_retriever.query_embedding",
)
pipeline.connect("bm25_retriever.documents", "joiner.documents")
pipeline.connect("embedding_retriever.documents", "joiner.documents")

query = "hybrid retrieval keyword semantic"
result = pipeline.run(
    data={
        "text_embedder": {"text": query},
        "bm25_retriever": {"query": query},
    },
    include_outputs_from={
        "bm25_retriever",
        "embedding_retriever",
    },
)

bm25_ids = [
    doc.id for doc in result["bm25_retriever"]["documents"]
]
embedding_ids = [
    doc.id for doc in result["embedding_retriever"]["documents"]
]
joined_documents = result["joiner"]["documents"]
joined_ids = [doc.id for doc in joined_documents]

assert "hybrid-search" in bm25_ids
assert "hybrid-search" in embedding_ids
assert joined_ids[0] == "hybrid-search"
assert "coffee-notes" not in joined_ids

print(f"Indexed documents: {document_store.count_documents()}")
print(f"BM25 candidates: {', '.join(bm25_ids)}")
print(f"Embedding candidates: {', '.join(embedding_ids)}")
print("Fused ranking:")
for position, document in enumerate(joined_documents, start=1):
    print(f"{position}. {document.id} | {document.score:.4f}")