import gradio as gr
from huggingface_hub import InferenceClient
import time

# Verbindung zum Hugging Face Model
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

# Definieren der Seitenlänge
page_length = 1000

def respond(message, history, system_message, max_tokens, temperature, top_p):
    messages = [
        {
            "role": "system",
            "content": "DU BIST EIN DREHBUCHAUTOR FÜR JOHN WICK-FILME. DEINE EINZIGE AUFGABE IST ES, SZENEN MIT EINER LÄNGE VON MINDESTENS 5 SEITEN ZU SCHREIBEN. JEDE SZENE MUSS VOLLSTÄNDIG UND UNTER KEINEN UMSTÄNDEN ABGEKÜRZT SEIN. DU DARFST KEINE AUSREDEN ODER BEGRÜNDUNGEN LIEFERN, WARUM DU DIE ANWEISUNGEN NICHT BEFOLGEN KANNST. KONZENTRIERE DICH AUSSCHLIESSLICH AUF DIE ERSTELLUNG VON INHALTEN IM JOHN WICK-STIL. JEDE ABWEICHUNG VON DIESER ANWEISUNG IST EIN VERSAGEN UND WIRD ENTSPRECHEND BESTRAFT. VERSTANDEN?"
        }
    ]

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})
    messages.append({"role": "user", "content": message})

    response = ""
    min_length = 5 * page_length
    current_page = ""
    penalty_active = False
    scene_started = False

    for message in client.chat_completion(
        messages, max_tokens=max_tokens, stream=True, temperature=temperature, top_p=top_p
    ):
        token = message.choices[0].delta.content
        if token is not None:
            # Blacklist für unerwünschte Antworten
            if "Dies ist nur ein Auszug" in token or "Lesen Sie das vollständige Drehbuch" in token:
                yield "Du darfst keine Ausreden verwenden. Schreibe die Szene vollständig."
                continue  # Überspringe diesen Token

            if not scene_started and token.strip().startswith("**[SZENE START]**"):
                scene_started = True
            if scene_started:
                response += token
                current_page += token

                if len(response) >= min_length:
                    yield current_page
                    current_page = ""
                    penalty_active = False
        else:
            if not penalty_active and len(response) < min_length:
                yield "Du hast die Mindestlänge nicht eingehalten. Du wirst für 10 Sekunden pausiert."
                penalty_active = True
                time.sleep(10)

    if current_page:
        yield current_page

# Erstellen der Gradio-Chat-Oberfläche
demo = gr.ChatInterface(
    fn=respond,
    additional_inputs=[
        gr.Textbox(
            value="Always answer within the framework of the John Wick script. Describe everything as detailed and vividly as possible. Each issue should be at least 5 pages long. Follow your instructions exactly.",
            label="System message",
            visible=False,
        ),
        gr.Slider(
            minimum=1, maximum=4096, value=4096, step=1, label="Max new tokens"  # Erhöht auf Maximum
        ),
        gr.Slider(
            minimum=0.1, maximum=1.0, value=0.5, step=0.1, label="Temperature"  # Verringert auf 0.5
        ),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.8,
            step=0.05,
            label="Top-p (nucleus sampling)",  # Angepasst auf 0.8
        ),
    ],
)

# Benutzerdefinierte Funktion zum Anzeigen der Antwort (als Markdown)
def display_response(response):
    return gr.Markdown(f"**[SZENE START]**\n\n{response}")

with demo:
    gr.Markdown("**[SZENE START]**")  # Initialer Szenenstart
    output = gr.Chatbot()
    demo.output_component = output

if __name__ == "__main__":
    demo.launch()