from huggingface_hub import InferenceClient
import os
import gradio as gr
import random
import time
from utils import *

# HF Inference Endpoints parameter
endpoint_url = "https://qrh4fv8e7x3fw9w3.us-east-1.aws.endpoints.huggingface.cloud"   

hf_token = os.getenv("TOKEN_HF") 

# Streaming Client
client = InferenceClient(endpoint_url, token=hf_token)
print ("Client ready")

########################################################################
#Hilfsfunktionen zur Darstellung inm Chatbot und File upload
# Chatbot demo with multimodal input (text, markdown, LaTeX, code blocks, image, audio, & video). Plus shows support for streaming text.
def add_text(history, text):
    history = history + [(text, None)]
    return history, gr.update(value="", interactive=False)


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

#wenn chat neu gestartet werden soll - Verlauf leeren und neue chatbotHF id
def reset_chat():
    return ""

#wenn 'Stop' Button geklickt, dann Message dazu und das Eingabe-Fenster leeren
def cancel_outputing():
    reset_textbox()
    return "Stop Done"

#ein Transfer-Zustand, bevor die Berechnung des Modells startet und die Antwort als Stream ausgegeben wird
#die letzte user-input wird aus dem textfeld gelöscht und schonmal in das Chatbot-Fenster geschrieben - darunter dann die Antwort als Stream
def reset_textbox_stream(message, chat_history):
    chat_history = chat_history + [[message, None]]
    #print(chat_history)
    #print("******************")
    return "", chat_history

#wenn button clear gedrückt wird, um chatbot zu leeren
def reset_state():
    return [], [], "Reset Done"

#user_input leeren
def reset_textbox():
    return gr.update(value=""),""

#message, wenn berechung per Stop Button unterbrochen wird
def cancel_outputing():
    return "Stop Done"



def user(user_message, history):
        return "", history + [[user_message, None]]

    

    
#Chat KI nutzen, um Text zu generieren...
def bot(user_message, history):
        # generation parameter
        gen_kwargs = dict(
            max_new_tokens=1024,
            top_k=30,
            top_p=0.9,
            temperature=0.2,
            repetition_penalty=1.02,
            stop_sequences=["\nUser:", "<|endoftext|>", "</s>"],
        )
        prompt = generate_prompt_with_history(user_message,history)
        print(prompt)
        stream = client.text_generation(prompt, stream=True, details=True, **gen_kwargs)
        print("++++++++++++++++++++++++++++stream++++++++++++++++++")

    
        history[-1][1] = ""
        # yield each generated token
        for r in stream:
            # skip special tokens
            if r.token.special:
                continue
            # stop if we encounter a stop sequence
            if r.token.text in gen_kwargs["stop_sequences"]:
                break
            # yield the generated token
            #print(r.token.text, end = "")
            history[-1][1] += r.token.text
            yield  history
            

                  

#######################################################################
#Darstellung mit Gradio

with open("custom.css", "r", encoding="utf-8") as f:
    customCSS = f.read()

with gr.Blocks() as demo:
    chatbot = gr.Chatbot([], elem_id="chatbot").style(height=550)
    status_display = gr.Markdown("Erfolg", elem_id="status_display")
    with gr.Row():
        with gr.Column(scale=0.85):
            txt = gr.Textbox(
                show_label=False,
                placeholder="Text eingeben und enter drücken oder ein Bild hochladen!",
            ).style(container=False)
        with gr.Column(scale=0.15, min_width=0):
            btn = gr.UploadButton("📁", file_types=["image", "video", "audio"])
    with gr.Row():
        with gr.Column(min_width=100, scale=0.5): 
            cancelBtn = gr.Button("Stoppen")
        with gr.Column(min_width=100, scale=0.5): 
            emptyBtn = gr.Button("🧹 Neuer Chat",)
        

    txt_msg = txt.submit(add_text, [chatbot, txt], [chatbot, txt], queue=False).then(
        bot, [txt, chatbot], chatbot
    )
    txt_msg.then(lambda: gr.update(interactive=True), None, [txt], queue=False)
    file_msg = btn.upload(add_file, [chatbot, btn], [chatbot], queue=False).then(
        bot, [txt, chatbot], chatbot
    )

    #Berechnung oder Ausgabe anhalten (kann danach fortgesetzt werden)
    cancelBtn.click(cancel_outputing, [], [status_display], cancels=[txt_msg]) 
    #cancelBtn.click(lambda: None, None, chatbotGr, queue=False)

    #neuer Chat - chatbotGr löschen und neue id für KI-Modell
    #wird chatbotGr gelöscht, dann auch der Kontext, der sonst ans Modell zur Berechnung gegeben wird
    reset_args = dict(
        fn=reset_chat, inputs=[], outputs=chatbot
    )
    #Listener, Wenn reset...
    emptyBtn.click(
        reset_state,
        outputs=chatbot,
        show_progress=True,
    )
    emptyBtn.click(**reset_args)

demo.queue()
demo.launch()