import os

from haystack import Document, Pipeline
from haystack.components.builders import ChatPromptBuilder
from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
from haystack.dataclasses import ChatMessage
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack_integrations.components.generators.ollama import OllamaChatGenerator


model = os.environ.get("OLLAMA_MODEL", "gpt-oss:20b")
url = os.environ.get("OLLAMA_URL", "http://localhost:11434")
question = "Which support team handles invoice corrections?"

document_store = InMemoryDocumentStore()
documents = [
    Document(
        content=(
            "Billing support handles invoice corrections and payment receipt questions."
        ),
        meta={"source": "billing-faq"},
    ),
    Document(
        content="Platform support handles login, password, and access token issues.",
        meta={"source": "access-faq"},
    ),
    Document(
        content="Shipping support handles delivery address changes before dispatch.",
        meta={"source": "shipping-faq"},
    ),
]
written = document_store.write_documents(documents)

retriever = InMemoryBM25Retriever(document_store=document_store, top_k=2)
retrieved = retriever.run(query=question)["documents"]

jinja_open = "{" + "{"
jinja_close = "}" + "}"

template = [
    ChatMessage.from_system(
        "Answer only from the supplied support FAQ context. "
        "If the answer is not in the context, say you do not know."
    ),
    ChatMessage.from_user(
        "Support FAQ context:\n"
        "{% for document in documents %}"
        "- " + jinja_open + " document.content " + jinja_close + "\n"
        "{% endfor %}\n"
        "Question: " + jinja_open + " question " + jinja_close + "\n"
        "Answer in one short sentence."
    ),
]

pipe = Pipeline()
pipe.add_component("retriever", retriever)
pipe.add_component(
    "prompt_builder",
    ChatPromptBuilder(template=template, required_variables=["documents", "question"]),
)
pipe.add_component(
    "llm",
    OllamaChatGenerator(
        model=model,
        url=url,
        generation_kwargs={"temperature": 0, "num_predict": 64},
        timeout=300,
        keep_alive="2m",
        think=False,
    ),
)
pipe.connect("retriever.documents", "prompt_builder.documents")
pipe.connect("prompt_builder.prompt", "llm.messages")

result = pipe.run(
    {
        "retriever": {"query": question},
        "prompt_builder": {"question": question},
    }
)
reply = result["llm"]["replies"][0].text.strip()

print(f"documents written: {written}")
print(f"model: {model}")
print(f"question: {question}")
print(f"top document source: {retrieved[0].meta['source']}")
print(f"top document: {retrieved[0].content}")
print(f"answer: {reply}")