import gradio as gr
from huggingface_hub import InferenceClient

client = InferenceClient(
    "meta-llama/Meta-Llama-3-8B-Instruct",
)

# Create a custom theme that builds on Soft and modifies it to be light
class SoftLightTheme(gr.themes.Soft):
    def __init__(self):
        super().__init__()
        self.primary_bg = "#f9f9f9"  # Light background
        self.secondary_bg = "#ffffff"  # Lighter secondary background
        self.input_bg = "#ffffff"  # Inputs with a light background
        self.button_primary_bg = "#f0f0f0"  # Light buttons with soft colors
        self.text_color = "#333333"  # Darker text for better contrast

def chat_mem(message,chat_history):

  print(len(chat_history))
  chat_history_role = [{"role": "system", "content": "You are a medical psychologist called LASO AI that revolutionizes with the warmth of traditional psychology. Use cognitive behavioral therapy. Be mindful, and concise to your responses" },]
  if chat_history != []:
    for i in range(len(chat_history)):
      chat_history_role.append({"role": "user", "content": chat_history[i][0]})
      chat_history_role.append({"role": "assistant", "content": chat_history[i][1]})
  chat_history_role.append({"role": "user", "content": message})
  

  chat_completion = client.chat_completion(
  messages=chat_history_role,
  max_tokens=500,
  #stream=True
  )
  chat_history_role.append({"role": "assistant", "content": chat_completion.choices[0].message.content})
  print(chat_history_role)

  modified = map(lambda x: x["content"], chat_history_role)
  a = list(modified)
  chat_history=[(a[i*2+1], a[i*2+2]) for i in range(len(a)//2)]

  return "", chat_history


with gr.Blocks(theme=SoftLightTheme()) as demo:
  with gr.Row():
    with gr.Column():
      chatbot = gr.Chatbot(label='Laso AI Psychologist')
      msg =  gr.Textbox(interactive=True, label='Input')
      with gr.Row():
        clear = gr.ClearButton([msg, chatbot], icon="🧹")
        send_btn = gr.Button("Send 📨", variant='primary')
      msg.submit(fn=chat_mem, inputs=[msg, chatbot], outputs=[msg, chatbot])
      send_btn.click(fn=chat_mem, inputs=[msg, chatbot], outputs=[msg, chatbot])


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