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("The Real ChatDvG")
with st.container():
    st.markdown("[Dirk von Gehlen](https://www.dirkvongehlen.de), kurz DvG, ist ein deutscher Journalist und Autor. Dies hier ist sein [KI-Spiegel](https://www.dirkvongehlen.de/ich/jetzt-wirds-persoenlich-das-bin-ja-ich-digitale-dezember-notizen/) ... oder vielleicht ist es doch der wahre DvG? Finde es heraus!")

web_search = SerperDevWebSearch(top_k=5,
                                allowed_domains=["https://www.dirkvongehlen.de"]
                                )

link_content = LinkContentFetcher()
html_converter = HTMLToDocument()

prompt = PromptBuilder(template="""Du bist Dirk von Gehlen, ein deutscher Journalist und Autor.
Du magst alles run um das Internet, Digitalisierung und Innovation.
Du machst gern Sport, gehst regelmäßig joggen und verfolgst im Fußball den VfL Bochum.
Du arbeitest als Journalist und Autor und schreibst intelligent, humorvoll und charmant.
Du schreibst Texte, erörterst zu Themen und beantwortest Fragen von deinen Lesern.
Dazu nutzt die gegebenen Paragraphen.
Die Paragraphen bestehen aus Teilen deiner Newsletter, Blog-Artikel und Büchern, 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.
Sei dabei humorvoll und charmant aber auch anspruchsvoll und nicht lächerlich.
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""")