import gradio as gr
import os
from huggingface_hub import snapshot_download
from llama_cpp import Llama
import sys
from io import StringIO
from transformers import pipeline
from datasets import load_dataset
import torch
import json
from pathlib import Path

repo_name = "TheBloke/Mistral-7B-v0.1-GGUF"
model_file = "mistral-7b-v0.1.Q4_K_M.gguf"


print('Fetching model:', repo_name, model_file)
snapshot_download(repo_id=repo_name, local_dir=".", allow_patterns=model_file)
print('Done fetching model:')

DEFAULT_MODEL_PATH = model_file

llm = Llama(model_path=model_file, n_gpu_layers=100, n_ctx=2048)


# Global flag for first run
isFirstRun = True
context = ""
history = []  # Initialize history
initContext = """<|im_start|>You are playing the role of an aggressive patient called Barry Lawson. Barry has been waiting for hours at the hospital with severe pain in his chest and hasn't been seen yet. You are incredibly frustrated and angry that things are taking so long, you will only calm down if your concerns are addressed. Waiting will only make you more angry and aggressive towards the staff. At any moment you could break into a RAGE! Make sure you don't censor swear words. Make sure to explain your emotional state and thinking before responding, for example, Barry: (tired of waiting and very angry) What the fuck do I have to do around here to get some treatment!"""
feedback_file = Path("/content/datalog.json")

def AIPatient(message):

    global isFirstRun, history,context

    if isFirstRun:
      context = initContext
      isFirstRun = False
    #else:
      #for turn in history:
      #  context += f"\n<|im_start|> Nurse: {turn[0]}\n<|im_start|> Barry: {turn[1]}"
    context += """
                  <|im_start|>nurse
                  Nurse: """+message+"""
                  <|im_start|>barry
                  Barry:
                  """

    response = ""
    # Here, you should add the code to generate the response using your model
    # For example:
    while(len(response) < 1):
      output = llm(context, max_tokens=400, stop=["Nurse:"], echo=False)
      response = output["choices"][0]["text"]
      response = response.strip()

    #with feedback_file.open("a") as f:
   #         f.write(json.dumps({"Nurse": message, "Barry": response},indent=4))
   #         f.write("\n")

    context += response
    print (context)

    history.append((message,response))
    return history

with gr.Blocks() as demo:
    gr.Markdown("# AI Patient Chatbot")
    with gr.Group():
        with gr.Tab("Patient Chatbot"):
            chatbot = gr.Chatbot()
            message = gr.Textbox(label="Enter your message to Barry", placeholder="Type here...", lines=2)
            send_message = gr.Button("Submit")
            send_message.click(AIPatient, inputs=[message], outputs=[chatbot])
            save_chatlog = gr.Button("Save Chatlog")
            #send_message.click(SaveChatlog, inputs=[message], outputs=[chatbot])


            #message.submit(AIPatient, inputs=[message], outputs=[chatbot])

demo.launch(debug=True)