import os
import gradio as gr
from langchain_community.vectorstores import Chroma
from transformers import RagTokenizer, RagSequenceForGeneration
from sentence_transformers import SentenceTransformer
from langchain import LLMChain, PromptTemplate
from langchain_community.llms import HuggingFacePipeline

#Konstanten
ANTI_BOT_PW = os.getenv("CORRECT_VALIDATE")

# Hugging Face Token direkt im Code setzen
hf_token = os.getenv("HF_READ")

# Initialisierung des Sentence-BERT Modells für die Embeddings
embedding_model = SentenceTransformer('all-MiniLM-L6-v2')

# Initialisierung von Tokenizer und RAG Modell mit Token
tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-nq", use_auth_token=hf_token)
model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-nq", use_auth_token=hf_token)

# Verbindung zur Chroma DB und Laden der Dokumente
chroma_db = Chroma(embedding_model=embedding_model, persist_directory = PATH_WORK + CHROMA_DIR)

# Erstellung eines HuggingFacePipeline LLM Modells
llm = HuggingFacePipeline(pipeline=model)


# Erstellen eines eigenen Retrievers mit Chroma DB und Embeddings
#retriever = chroma_db.as_retriever()

# Erstellung der RAG-Kette mit dem benutzerdefinierten Retriever
#rag_chain = RagChain(model=model, retriever=retriever, tokenizer=tokenizer, vectorstore=chroma_db)
#############################################


def document_retrieval_chroma2(): 
    #HF embeddings -----------------------------------
    #Alternative Embedding - für Vektorstore, um Ähnlichkeitsvektoren zu erzeugen - die ...InstructEmbedding ist sehr rechenaufwendig
    embeddings = HuggingFaceInstructEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2", model_kwargs={"device": "cpu"})
    #etwas weniger rechenaufwendig:
    #embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2", model_kwargs={"device": "cpu"}, encode_kwargs={'normalize_embeddings': False})
    #oder einfach ohne Langchain:
    #embeddings = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")

    #ChromaDb um die embedings zu speichern
    db = Chroma(embedding_function = embeddings, persist_directory = PATH_WORK + CHROMA_DIR)
    print ("Chroma DB bereit ...................")
    
    return db 
	
	
	
def get_rag_response(question):
    # Abfrage der relevanten Dokumente aus Chroma DB
    docs = chroma_db.search(question, top_k=5)
    passages = [doc['text'] for doc in docs]
    links = [doc.get('url', 'No URL available') for doc in docs]

    # Generieren der Antwort
    prompt_template = PromptTemplate(input_variables=["context", "question"], template="{context}\n\n{question}")
    prompt = prompt_template(context=" ".join(passages), question=question)
    answer = llm(prompt)
    
    # Zusammenstellen der Ausgabe
    response = {
        "answer": answer,
        "documents": [{"link": link, "passage": passage} for link, passage in zip(links, passages)]
    }
    
    return response

	
	
# Funktion, die für den Chatbot genutzt wird
def chatbot_response(user_input, chat_history=[]):
    response = get_rag_response(user_input)
    answer = response['answer']
    documents = response['documents']
    
    doc_links = "\n\n".join([f"Link: {doc['link']}\nPassage: {doc['passage']}" for doc in documents])
    
    bot_response = f"{answer}\n\nRelevant Documents:\n{doc_links}"
    
    chat_history.append((user_input, bot_response))
    return chat_history, chat_history
	
	
#############################
#GUI.........
def user (user_input, history):
	return "", history + [[user_input, None]]

with gr.Blocks() as chatbot:
	chat_interface = gr.Chatbot()
	msg = gr.Textbox()
	clear = gr.Button("Löschen")
	
	#Buttons listener
	msg.submit(user, [msg, chat_interface], [msg, chat_interface], queue = False). then(chatbot_response, [msg, chat_interface], [chat_interface, chat_interface])
	
	clear.click(lambda: None, None, chat_interface, queue=False)
	
chatbot.launch()