import os
import openai
import gradio as gr

openai.api_key = os.environ['key']

model_name = "ft:gpt-3.5-turbo-0613:metric-space-ug:rechtsberatung2:7zmeC6ps"
prompt0 = "Liegt ein rechtlich relevanter Sachverhalt mit einer passenden Frage vor?"
prompt1 = "Du stellst Fragen, um fehlende Informationen für eine anwaltliche Erstberatung zu klären. Stelle bis zu 10 Fragen für die Klärung der Sachlage und antworte mit „Danke für die Antworten.“, sobald die Sachlage klar ist. Vermeide Fragen zu stellen, die schon benantwortet sind"
prompt2 = "Schreibe den Sachverhalt auf Basis der folgenden Informationen zusammen."
prompt3 = "Schreibe eine anwaltliche Erstberatung auf Basis des geschilderten Sachverhalts. Vermeide Konjuktive, formuliere in gutem Deutsch und kurzen Sätzen."

def respond_prompt0(Selbstauskunft, Kernfrage, chat):
  # preprocess data

  # openai api call prompt 0
  ########
  completion = openai.ChatCompletion.create(
  model=model_name,
  messages=[
    {"role": "system", "content": prompt0},
    {"role": "user", "content": Selbstauskunft + "\n" + Kernfrage}
  ])
  output = completion.choices[0].message.content
  ########

  out = "Kann ich dir mit deinem Anliegen helfen?\n" + output
  if "ja" in output[:4].lower():
    out += "\n Ich werde nun noch ein paar Fragen stellen, um den Sachverhalt genauer zu verstehen" #

    # openai api call prompt 1
    ########
    completion = openai.ChatCompletion.create(
    model=model_name,
    messages=[
      {"role": "system", "content": prompt1},
      {"role": "user", "content": Selbstauskunft + "\n" + Kernfrage}
    ])
    output_prompt1 = completion.choices[0].message.content
    ########

    chat = [(output_prompt1, "")]
    return out, gr.update(value=chat, visible=True), gr.update(value="", visible=True)

  else:
    return out + ". \nLeider kann ich dir nicht weiterhelfen.", chat, ""


def respond_prompt1(Selbstauskunft, Kernfrage, chat, text):
  chat[-1][1] = text
  # openai api call prompt 1
  ########
  messages=[
      {"role": "system", "content": prompt1},
      {"role": "user", "content": Selbstauskunft + "\n" + Kernfrage}
    ]

  for i in range(len(chat)):
    messages.append({ "role": "assistant", "content": chat[i][0]})
    messages.append({ "role": "user", "content": chat[i][1] })

  completion = openai.ChatCompletion.create(
  model=model_name,
  messages=messages)

  output_prompt1 = completion.choices[0].message.content
  ########


  if ("<ENDOFQA>" in output_prompt1) or len(chat) > 10:
    sachverhalt = respond_prompt2(Selbstauskunft, Kernfrage, chat).replace("Bürger", "Mandant")
    chat.append(("Danke für deine Antworten. Hier findest du eine aufarbeitete Zusamenfassung des Sachverhalts. Bitte prüfe dieee und ergänze ggf. noch etwas, bevor ich zu meiner Einschätzung komme.", ""))

    return chat,  gr.update(value="", visible=False),  gr.update(value=sachverhalt, visible=True), gr.update(value="Weiter zur juristischen Einschätzung", visible=True)

  chat.append((output_prompt1, ""))

  return chat, "", "", ""

def respond_prompt2(Selbstauskunft, Kernfrage, chat):
  # openai api call prompt 2
  ########
  content = Selbstauskunft + "\n" + Kernfrage

  for i in range(len(chat)-1):
    content += "\n\n" + chat[i][0] + "\n"
    content +=  chat[i][1]

  completion = openai.ChatCompletion.create(
  model=model_name,
  messages=[
  { "role": "system", "content": prompt2 },
  { "role": "user", "content": content}])

  output_prompt2 = completion.choices[0].message.content
  ########

  return output_prompt2

def respond_prompt3(Selbstauskunft, Kernfrage, chat, Sachverhalt):
  # openai api call prompt 3
  ########
  content = Selbstauskunft + "\n" + Kernfrage

  for i in range(len(chat)-1):
    content += "\n\n" + chat[i][0] + "\n"
    content +=  chat[i][1]

  completion = openai.ChatCompletion.create(
  model=model_name,
  messages=[
  { "role": "system", "content": prompt3 },
  { "role": "user", "content": content +  "\nSachverhalt: " + Sachverhalt}])


  output_prompt3 =completion.choices[0].message.content + "\n\n(Dies ist eine maschinell erstellte Einschätzung einer experimentellen KI-Software basierend auf Informationen im Internet und keine Anwaltliche Erstberatung eines Anwalts. Interpretieren Sie das Ergebnis daher mit Vorsicht.)"
  ########

  return gr.update(value=output_prompt3, visible=True), gr.update(visible=False)

with gr.Blocks() as demo:
    Title = gr.title("Juristische Ersteinschätzung einer KI.")
    desc = gr.description("Disclaimer: Dies ist keine Anwaltliche Erstberatung, sondern ein Forschungsprojekt.")
    Selbstauskunft = gr.Textbox("", label="Erzähl mir, was passiert ist.")
    Kernfrage = gr.Textbox("", label="Stelle eine Frage, die ich dir beantworten soll.")
    Selbstauskunft = gr.Textbox("", label="Erzähl mir, was passiert ist.")
    Kernfrage = gr.Textbox("", label="Stelle eine Frage, die ich dir beantworten soll.")

    button_prompt1 =  gr.Button(value="Weiter")

    label_prompt0 = gr.Label()


    chat = gr.Chatbot(visible=False)
    msg = gr.Textbox(visible=False, label="Antwort")
    Sachverhalt = gr.Textbox(visible=False, label = "Zusammenfassung")
    Erstberatung = gr.Textbox(visible=False, label="Rechtliche Einschätzung")
    button_erstberatung = gr.Button(value="Weiter zur juristischen Einschätzung", visible=False)

    button_prompt1.click(respond_prompt0, [Selbstauskunft, Kernfrage, chat], outputs=[label_prompt0, chat, msg])

    msg.submit(respond_prompt1, [Selbstauskunft, Kernfrage, chat, msg], [chat, msg, Sachverhalt, button_erstberatung])

    button_erstberatung.click(respond_prompt3, [Selbstauskunft, Kernfrage, chat, Sachverhalt], [Erstberatung, button_erstberatung])


demo.launch()