import requests
from huggingface_hub import InferenceClient, login
from transformers import AutoTokenizer
from langchain.chat_models import ChatOpenAI
import os, sys, json
import gradio as gr
from langchain.evaluation import load_evaluator
from pprint import pprint as print
import time

from langchain.chains import LLMChain, RetrievalQA
from langchain.chat_models import ChatOpenAI
from langchain.retrievers.self_query.base import SelfQueryRetriever
from langchain.document_loaders import PyPDFLoader, WebBaseLoader, UnstructuredWordDocumentLoader, DirectoryLoader
from langchain.document_loaders.blob_loaders.youtube_audio import YoutubeAudioLoader
from langchain.document_loaders.generic import GenericLoader
from langchain.document_loaders.parsers import OpenAIWhisperParser
from langchain.schema import AIMessage, HumanMessage
from langchain.embeddings import HuggingFaceInstructEmbeddings, HuggingFaceEmbeddings, HuggingFaceBgeEmbeddings, HuggingFaceInferenceAPIEmbeddings
from langchain.prompts import PromptTemplate
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.retrievers.self_query.base import SelfQueryRetriever
from langchain.vectorstores import Chroma
from chromadb.errors import InvalidDimensionException


from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv())


# access token with permission to access the model and PRO subscription
HUGGINGFACEHUB_API_TOKEN = os.getenv("HF_ACCESS_READ")
os.environ["HUGGINGFACEHUB_API_TOKEN"] = HUGGINGFACEHUB_API_TOKEN
#login(token=os.environ["HF_ACCESS_READ"])

OAI_API_KEY=os.getenv("OPENAI_API_KEY")



#################################################
#Prompt Zusätze
#################################################
template = """Antworte in deutsch, wenn es nicht explizit anders gefordert wird. Wenn du die Antwort nicht kennst, antworte einfach, dass du es nicht weißt. Versuche nicht, die Antwort zu erfinden oder aufzumocken. Halte die Antwort kurz aber exakt. Wenn du die Antwort in den Dokumenten nicht findest, schaue in deinen eigenen Daten nach, mache das aber in der Antwort deutlich."""

llm_template = "Beantworte die Frage am Ende. " + template + "Frage: {question} Hilfreiche Antwort: "
rag_template = "Nutze die folgenden Kontext Teile, um die Frage zu beantworten am Ende. " + template + "{context} Frage: {question} Hilfreiche Antwort: "


#################################################
# Konstanten
#RAG: Pfad, wo Docs/Bilder/Filme abgelegt werden können - lokal, also hier im HF Space (sonst auf eigenem Rechner)
#################################################
PATH_WORK = "."
CHROMA_DIR  = "/chroma"
YOUTUBE_DIR = "/youtube"

###############################################
#URLs zu Dokumenten oder andere Inhalte, die einbezogen werden sollen
PDF_URL       = "https://arxiv.org/pdf/2303.08774.pdf"
WEB_URL       = "https://openai.com/research/gpt-4"
YOUTUBE_URL_1 = "https://www.youtube.com/watch?v=--khbXchTeE"
YOUTUBE_URL_2 = "https://www.youtube.com/watch?v=hdhZwyf24mE"
#YOUTUBE_URL_3 = "https://www.youtube.com/watch?v=vw-KWfKwvTQ"


###############################################
#globale Variablen
##############################################
#nur bei ersten Anfrage splitten der Dokumente - um die Vektordatenbank entsprechend zu füllen
splittet = False


##############################################
# inference client
##############################################
print ("Inf.Client")
#client = InferenceClient("https://api-inference.huggingface.co/models/meta-llama/Llama-2-70b-chat-hf")
#client = InferenceClient("https://ybdhvwle4ksrawzo.eu-west-1.aws.endpoints.huggingface.cloud")
#Inference mit Authorisation:
#API_URL = "https://api-inference.huggingface.co/models/HuggingFaceH4/zephyr-7b-beta"
HEADERS = {"Authorization": f"Bearer {HUGGINGFACEHUB_API_TOKEN}"}

##############################################
# tokenizer for generating prompt
##############################################
print ("Tokenizer") 
#tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-70b-chat-hf")
#tokenizer = AutoTokenizer.from_pretrained("HuggingFaceH4/zephyr-7b-beta")
#tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.1")
#tokenizer =  AutoTokenizer.from_pretrained("mistralai/Mixtral-8x7B-Instruct-v0.1")

##############################################
# Zum Testen:
#list of models available
#client = InferenceClient()
#print("List of models ......................:")
#print(client.list_deployed_models("text-generation-inference"))

#angezeigt am 17.12.2023:
#{'text-generation': ['bigcode/starcoder','bigscience/bloom','codellama/CodeLlama-13b-hf','codellama/CodeLlama-34b-Instruct-hf','HuggingFaceH4/zephyr-7b-beta','HuggingFaceM4/idefics-80b-instruct', 'meta-llama/Llama-2-70b-chat-hf',
                     #'mistralai/Mistral-7B-Instruct-v0.1','mistralai/Mistral-7B-v0.1', 'OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5','openchat/openchat_3.5','TheBloke/vicuna-7B-v1.5-GPTQ','tiiuae/falcon-180B-chat','tiiuae/falcon-7b',
                     #'tiiuae/falcon-7b-instruct'],'text2text-generation': ['google/flan-t5-xxl']}




#################################################
#################################################
#################################################
#Funktionen zur Verarbeitung
################################################
def add_text(history, text):
    history = history + [(text, None)]
    return history, gr.Textbox(value="", interactive=False)


def add_file(history, file):
    history = history + [((file.name,), None)]
    return history

################################################
################################################
# Für den Vektorstore...
# Funktion, um für einen best. File-typ ein directory-loader zu definieren
def create_directory_loader(file_type, directory_path):
    #verschiedene Dokument loaders:
    loaders = {
        '.pdf': PyPDFLoader,
        '.word': UnstructuredWordDocumentLoader,
    }
    return DirectoryLoader(
        path=directory_path,
        glob=f"**/*{file_type}",
        loader_cls=loaders[file_type],
    )
    
#die Inhalte splitten, um in Vektordatenbank entsprechend zu laden als Splits
def document_loading_splitting():
    global splittet
    ##############################
    # Document loading
    docs = []
    
    # kreiere einen DirectoryLoader für jeden file type
    pdf_loader = create_directory_loader('.pdf', './chroma/pdf')
    word_loader = create_directory_loader('.word', './chroma/word')  
    
    # Laden der files
    pdf_documents = pdf_loader.load()
    word_documents = word_loader.load()

    #alle zusammen in docs (s.o.)...
    docs.extend(pdf_documents)
    docs.extend(word_documents)

    #andere loader  - für URLs zu Web, Video, PDF im Web...
    # Load PDF
    loader = PyPDFLoader(PDF_URL)
    docs.extend(loader.load())
    # Load Web
    loader = WebBaseLoader(WEB_URL)
    docs.extend(loader.load())
    # Load YouTube
    loader = GenericLoader(YoutubeAudioLoader([YOUTUBE_URL_1,YOUTUBE_URL_2], PATH_WORK + YOUTUBE_DIR), OpenAIWhisperParser())
    docs.extend(loader.load())
    
    ################################
    # Vektorstore Vorbereitung: Document splitting
    text_splitter = RecursiveCharacterTextSplitter(chunk_overlap = 150, chunk_size = 1500)
    splits = text_splitter.split_documents(docs)
    
    #nur bei erster Anfrage mit "choma" wird gesplittet...
    splittet = True
    
    return splits

#Vektorstore anlegen...
#Chroma DB die splits ablegen - vektorisiert...
def document_storage_chroma(splits):
    #OpenAi embeddings----------------------------------
    Chroma.from_documents(documents = splits, embedding = OpenAIEmbeddings(disallowed_special = ()),  persist_directory = PATH_WORK + CHROMA_DIR)  

    #HF embeddings--------------------------------------
    #Chroma.from_documents(documents = splits, embedding = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2", model_kwargs={"device": "cpu"}, encode_kwargs={'normalize_embeddings': False}),  persist_directory = PATH_WORK + CHROMA_DIR)  
    

#Vektorstore vorbereiten...
#dokumente in chroma db vektorisiert ablegen können - die Db vorbereiten daüfur
def document_retrieval_chroma(): 
    #OpenAI embeddings -------------------------------
    embeddings = OpenAIEmbeddings()

    #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 



###############################################
#Langchain anlegen

#langchain nutzen, um prompt an LLM zu leiten - llm und prompt sind austauschbar
#prompt ohne RAG!!!
def llm_chain(prompt):
    llm_template = "Beantworte die Frage am Ende. " + template + "Frage: " + prompt
    return llm_template

#langchain nutzen, um prompt an llm zu leiten, aber vorher in der VektorDB suchen, um passende splits zum Prompt hinzuzufügen
#prompt mit RAG!!!
def rag_chain(prompt, db, k=3):
    rag_template = "Nutze die folgenden Kontext Teile am Ende, um die Frage zu beantworten . " + template + "Frage: " + prompt + "Kontext Teile: "
    retrieved_chunks  = db.similarity_search(prompt)

    neu_prompt = rag_template 
    for i, chunk in enumerate(retrieved_chunks):
        neu_prompt += f"{i+1}. {chunk}\n"

    return neu_prompt



###################################################
#Prompts mit History erzeugen für verschiednee Modelle
###################################################
#Funktion, die einen Prompt mit der history zusammen erzeugt - allgemein
def generate_prompt_with_history(text, history, max_length=4048):
    #prompt = "The following is a conversation between a human and an AI assistant named Baize (named after a mythical creature in Chinese folklore). Baize is an open-source AI assistant developed by UCSD and Sun Yat-Sen University. The human and the AI assistant take turns chatting. Human statements start with [|Human|] and AI assistant statements start with [|AI|]. The AI assistant always provides responses in as much detail as possible, and in Markdown format. The AI assistant always declines to engage with topics, questions and instructions related to unethical, controversial, or sensitive issues. Complete the transcript in exactly that format.\n[|Human|]Hello!\n[|AI|]Hi!"   
    #prompt = "Das folgende ist eine Unterhaltung in deutsch zwischen einem Menschen und einem KI-Assistenten, der Baize genannt wird. Baize ist ein open-source KI-Assistent, der von UCSD entwickelt wurde. Der Mensch und der KI-Assistent chatten abwechselnd miteinander in deutsch. Die Antworten des KI Assistenten sind immer so ausführlich wie möglich und in Markdown Schreibweise und in deutscher Sprache. Wenn nötig übersetzt er sie ins Deutsche. Die Antworten des KI-Assistenten vermeiden Themen und Antworten zu unethischen, kontroversen oder sensiblen Themen. Die Antworten sind immer sehr höflich formuliert..\n[|Human|]Hallo!\n[|AI|]Hi!"   
    prompt=""
    history = ["\n{}\n{}".format(x[0],x[1]) for x in history]
    history.append("\n{}\n".format(text))
    history_text = ""
    flag = False
    for x in history[::-1]:
            history_text = x + history_text
            flag = True
    print ("Prompt: ..........................")
    print(prompt+history_text)
    if flag:
        return  prompt+history_text
    else:
        return None

#Prompt und History für OPenAi Schnittstelle
def generate_prompt_with_history_openai(prompt, history):
    history_openai_format = []
    for human, assistant in history:
        history_openai_format.append({"role": "user", "content": human })
        history_openai_format.append({"role": "assistant", "content":assistant})

    history_openai_format.append({"role": "user", "content": prompt})
    return history_openai_format
        

##############################################
##############################################
##############################################
# generate function
##############################################
def generate(text, history, rag_option, model_option,  k=3, temperature=0.5, max_new_tokens=4048, top_p=0.6, repetition_penalty=1.3,):
    #mit RAG
    if (rag_option is None):
        raise gr.Error("Retrieval Augmented Generation ist erforderlich.")
    if (text == ""):
        raise gr.Error("Prompt ist erforderlich.")

    try:
        if (model_option == "HuggingFace1"):
            #Anfrage an InferenceEndpoint1 ----------------------------
            API_URL = "https://api-inference.huggingface.co/models/HuggingFaceH4/zephyr-7b-beta"
            print("HF1")
        else:
            API_URL = "https://api-inference.huggingface.co/models/OpenAssistant/oasst-sft-4-pythia-12b-epoch-3.5"
            print("HF2")
            
        if (rag_option == "An"):
            #muss nur einmal ausgeführt werden... 
            if not splittet:
                splits = document_loading_splitting()
                document_storage_chroma(splits)
            db = document_retrieval_chroma()
            #mit RAG:
            neu_text_mit_chunks = rag_chain(text, db, k)
            #für Chat LLM:
            #prompt = generate_prompt_with_history_openai(neu_text_mit_chunks, history)
            #als reiner prompt:
            prompt = generate_prompt_with_history(neu_text_mit_chunks, history)
        else:
            #für Chat LLM:
            #prompt = generate_prompt_with_history_openai(text, history)
            #als reiner prompt:
            prompt = generate_prompt_with_history(text, history)
            
        #Anfrage an Modell (mit RAG: mit chunks aus Vektorstore, ohne: nur promt und history)
        #payload = tokenizer.apply_chat_template([{"role":"user","content":prompt}],tokenize=False)
        #Für LLAMA:
        #payload = tokenizer.apply_chat_template(prompt,tokenize=False)
        #result = client.text_generation(payload, do_sample=True,return_full_text=False, max_new_tokens=2048,top_p=0.9,temperature=0.6,)    
        #inference allg:
        data = {
            "inputs": prompt,
            "options": {"max_new_tokens": max_new_tokens},
        }
        response= requests.post(API_URL, headers=HEADERS, json=data)
        result = response.json()
        print("result:------------------")
        chatbot_response = result[0]['generated_text']
        print(chatbot_response)
        print("anzahl tokens gesamt antwort:------------------")
        print (len(chatbot_response.split()))
    except Exception as e:
        raise gr.Error(e)

    chatbot_message = chatbot_response[len(prompt):].strip()
    #return chatbot_message 
    
    #Antwort als Stream ausgeben... 
    for i in range(len(chatbot_message)):
        time.sleep(0.03)
        yield  chatbot_message[: i+1]
    
    
    
    
    #zum Evaluieren:
    # custom eli5 criteria
    #custom_criterion = {"eli5": "Is the output explained in a way that a 5 yeard old would unterstand it?"}

    #eval_result = evaluator.evaluate_strings(prediction=res.strip(), input=text, criteria=custom_criterion, requires_reference=True)
    #print ("eval_result:............    ")
    #print(eval_result)
    #return res.strip()

########################################
#Evaluation
########################################
evaluation_llm = ChatOpenAI(model="gpt-4")

# create evaluator
evaluator = load_evaluator("criteria", criteria="conciseness", llm=evaluation_llm)


################################################
#GUI
###############################################
#Beschreibung oben in GUI
################################################
print ("Start GUI")

description = """<strong>Information:</strong> Hier wird ein <strong>Large Language Model (LLM)</strong> mit 
                 <strong>Retrieval Augmented Generation (RAG)</strong> auf <strong>externen Daten</strong> verwendet.\n\n
                 """
css = """.toast-wrap { display: none !important } """
examples=[['Was ist ChtGPT-4?'],['schreibe ein Python Programm, dass die GPT-4 API aufruft.']]

additional_inputs = [
    gr.Radio(["Aus", "An"], label="RAG - LI Erweiterungen", value = "Aus"),
    gr.Radio(["HuggingFace1", "HuggingFace2"], label="Modellauswahl", value = "HuggingFace1"),
    gr.Slider(label="Anzahl Vergleichsdokumente", value=3, minimum=1, maximum=10, step=1, interactive=True, info="wie viele Dokumententeile aus dem Vektorstore an den prompt gehängt werden", visible=True),
    gr.Slider(label="Temperature", value=0.65, minimum=0.0, maximum=1.0, step=0.05, interactive=True, info="Höhere Werte erzeugen diversere Antworten", visible=True),
    gr.Slider(label="Max new tokens", value=1024, minimum=0, maximum=4096, step=64, interactive=True, info="Maximale Anzahl neuer Tokens", visible=True),
    gr.Slider(label="Top-p (nucleus sampling)", value=0.6, minimum=0.0, maximum=1, step=0.05, interactive=True, info="Höhere Werte verwenden auch Tokens mit niedrigerer Wahrscheinlichkeit.", visible=True),
    gr.Slider(label="Repetition penalty", value=1.2, minimum=1.0, maximum=2.0, step=0.05, interactive=True, info="Strafe für wiederholte Tokens", visible=True)
]

chatbot_stream = gr.Chatbot()

chat_interface_stream = gr.ChatInterface(fn=generate, 
                additional_inputs = additional_inputs,
                title = "ChatGPT vom LI",
                theme="soft",
                chatbot=chatbot_stream,
                retry_btn="🔄 Wiederholen",
                undo_btn="↩️ Letztes löschen",
                clear_btn="🗑️ Verlauf löschen",
                submit_btn = "Abschicken",
                description = description) 

with gr.Blocks() as demo:
    with gr.Tab("Chatbot"):
        #chatbot_stream.like(vote, None, None)
        chat_interface_stream.queue().launch()