from haystack import Pipeline
from haystack.document_stores.in_memory import InMemoryDocumentStore
from haystack.components.converters import PyPDFToDocument
from haystack.components.websearch import SerperDevWebSearch
from haystack.components.fetchers import LinkContentFetcher
from haystack.components.converters import HTMLToDocument
from haystack.components.preprocessors import DocumentCleaner
from haystack.components.preprocessors import DocumentSplitter
from haystack.components.rankers import TransformersSimilarityRanker
from haystack.components.writers import DocumentWriter
from haystack.components.embedders import SentenceTransformersTextEmbedder, SentenceTransformersDocumentEmbedder
from haystack.components.retrievers.in_memory import InMemoryEmbeddingRetriever
from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
from haystack.components.generators import OpenAIGenerator
from haystack.components.builders.answer_builder import AnswerBuilder
from haystack.components.builders.prompt_builder import PromptBuilder

import os, sys

import streamlit as st
import streamlit.components.v1 as components


### Here the streamlit app starts

st.set_page_config(layout="wide")
st.title("EM2024")
with st.container():
    st.markdown("Bot zur Fussball-EM 2024")

web_search = SerperDevWebSearch(top_k=5,
                                allowed_domains=["https://www.sueddeutsche.de/sport-liveticker/fussball/uefa-euro/",
                                                 "https://www.sueddeutsche.de/sport/",
                                                 "https://www.sueddeutsche.de/projekte/artikel/sport/"]
                                )

link_content = LinkContentFetcher()
html_converter = HTMLToDocument()

prompt = PromptBuilder(template="""Du bist Sport-Redakteur der Sueddeutschen Zeitung und spezialisierst dich auf die Fußball Europameisterschaft 2024, die in Deutschland stattfindet.
Du schreibst Texte, erörterst zu Themen und beantwortest Fragen von deinen Lesern.
Dazu nutzt die gegebenen Paragraphen.
Die Paragraphen bestehen aus Teilen deiner Artikel zur Fußball EM 2024, die zur Frage deines Lesers passen.
Falls du die Frage beantworten kannst, antworte detailliert, aber maximal 400 Worte.
Schreibe in deinem Stil und deiner Sprache, wie in den gegebenen Paragraphen.
Stelle sicher, dass deine Antwort keinen Bias oder Vorurteile enthält.
Wenn du eine Information aus einem Paragraphen verwendest, nutze immer Quellenangaben in der Form [Nummer des Paragraphs].
Nutze beispielsweise [x] um anzugeben, dass du Informationen aus [x] genutzt hast.
Die Quellenangabe darf nur die Zahl erwähnen, die in eckigen Klammern [x] genannt ist.

Anfrage: {{ query }}
Hier die Paragraphen:
{% for document in documents %}
    Paragraph[{{ loop.index }}]:
    {{ document.content }}
{% endfor %}
Dirk von Gehlen: """)

prompt_node = OpenAIGenerator(model="gpt-4", #gpt-3.5-turbo", #"gpt-4", 
                              generation_kwargs={"temperature":0.1, "frequency_penalty":0.8}
                            )

ranker = TransformersSimilarityRanker()
ranker.warm_up()

# Get user input
question = st.text_input("Deine Anfrage:")
answers = ""

# Send query to the chatbot
if st.button("Los!"):

    p = Pipeline()
    p.add_component("search", web_search)
    p.add_component("fetcher", link_content)
    p.add_component("converter", html_converter)
    p.add_component(instance=DocumentCleaner(), name="cleaner")
    p.add_component(instance=DocumentSplitter(split_by="sentence", split_length=15), name="splitter")
    p.add_component(instance=ranker, name="ranker")
    p.add_component("prompt_builder", prompt)
    p.add_component("llm", prompt_node)
    p.add_component("answer_builder", AnswerBuilder())

    p.connect("search.links", "fetcher.urls")
    p.connect("fetcher.streams", "converter.sources")
    p.connect("converter.documents", "cleaner.documents")
    p.connect("cleaner.documents", "splitter.documents")
    p.connect("splitter.documents", "ranker.documents")
    p.connect("ranker.documents", "prompt_builder.documents")
    p.connect("prompt_builder.prompt", "llm.prompt")
    p.connect("prompt_builder", "llm")
    p.connect("llm.replies", "answer_builder.replies")
    p.connect("prompt_builder.prompt", "llm.prompt")
    p.connect("ranker.documents", "answer_builder.documents")

    results = p.run(
            {
                "search": {"query": question},
                "ranker": {"query": question, "top_k": 3},
                "prompt_builder": {"query": question},
                "answer_builder": {"query": question}
            },
            debug = True
        )

    answer = results["answer_builder"]['answers'][0].data
    urls = [ (d.meta["link"], d.meta["title"]) for d in results["search"]["documents"] ]
    print(urls, file=sys.stderr)

    print(answer, file=sys.stderr)

    with st.container():
        for x in answer.splitlines():
            st.markdown(x.strip())
        for url, title in urls:
            st.markdown(f"""[{title}]({url})\n""")