import os
from openai import OpenAI
from dotenv import load_dotenv
import gradio as gr

load_dotenv()
API_KEY = os.getenv("OPENAI_API_KEY")
openai = OpenAI(api_key=API_KEY)

create_msg = lambda x, y: {"role": x, "content": y}

SYSTEM_PROMPT = create_msg(
    "system",
    """You are a helpful mental health chatbot, please answer with care. If you don't know the answer, respond 'Sorry, I don't know the answer to this question.'. If the question is too complex, respond 'Kindly, consult a psychiatrist for further queries.'.""".strip(),
)


def predict(message, history):
    history_openai_format = []
    history_openai_format.append(SYSTEM_PROMPT)
    for human, assistant in history:
        history_openai_format.append({"role": "user", "content": human})
        history_openai_format.append({"role": "assistant", "content": assistant})
    history_openai_format.append({"role": "user", "content": message})

    response = openai.chat.completions.create(
        model="ft:gpt-3.5-turbo-0613:personal::8qTSdtfs", messages=history_openai_format, temperature=0.35, stream=True
    )

    partial_message = ""
    for chunk in response:
        if chunk.choices[0].delta.content is not None:
            partial_message = partial_message + chunk.choices[0].delta.content
            yield partial_message


gr.ChatInterface(
    fn=predict,
    title="Mental Health Chatbot by Jayda Hunte",
    description="This is a mental health companion created to assist persons who need an outlet to express their raw feelings without the fear of being judged. Sometimes persons don't have a direct outlet they trust to talk about certain aspects of their lives, so this chatbot will serve as a companion in those times of need. This is a pure companion who remembers your name throughout each session and provides an outlet for you to vent and get actionable advice.",
    examples=[
        "Hi! My name is John and I haven't been feeling myself lately. A lot has been going and I just don't know how to process it all. I feel like I'm losing my mind.",
        "I have been feeling inadequate lately - I feel like I'm not good enough. Everyone around me is doing amazing things and I'm just here, stuck in the same place. I don't know what to do.",
        "I have been feeling very anxious lately. I haven't been myself and it is really bothering me. What steps can I take to get out of this rut?",
    ],
).launch()