from nltk.metrics.distance import edit_distance
import openai
import os
import questionnaires

initial_message = "Hi! I'm Brainy, a virtual assistant who can help you through the mental health screening process. Can you start by telling me a bit about the problems you are experiencing?"
services = "https://www.canada.ca/en/public-health/topics/mental-health-wellness.html"

openai.api_key = os.getenv("OPENAI_API_KEY")

available_questionnaires = list(questionnaires.topics)

def listify(options: list, end_and=True, lower=True, use_format=None):
    options = options.copy()
    if end_and is True and len(options) > 1: options[-1] = "and " + options[-1]
    if lower is True: options = list(map(str.lower, options))
    if use_format: list(map(use_format.format, options))
    return ", ".join(options) if len(options) > 2 else " ".join(options)

def content_summarizer():
    """
    Summarize previous messages for new prompts to pass as context.
    """
    pass

def response_aligner(options:list, response:str): 
    """
    Takes a series of options, and a response as given by Openai.
    Align the response to the nearest option (edit distance).
    This is necessary as sometimes openai will return a response slightly different to the ones in prompt.
    """
    distances = [(option, edit_distance(option, response)) for option in options]
    distances = sorted(distances, key=lambda item: item[1])
    return distances[0][0]

def assess_response(question: str, response:str, options: list, uses_scale=True, context="a psychological evaluation"):
    context_prompt = "" if context=="" else f"The context of this discussion is {context}. "
    if uses_scale is True:
        prompt = context_prompt + f"""
        Given the question '{question}' and this response to the question '{response}', rate the response on the following scale: {listify(options, end_and=False, lower=False, use_format="'{}'")}. 
        Return only the corresponding scale item and nothing else.
        """
    else:
        prompt = context_prompt + f"""
        Given the question '{question}' and this response to the question '{response}', which of the following is the response most similar to? {listify(options, end_and=False, lower=False, use_format="'{}'")}. 
        Return only the most similar item and nothing else.
        """
    response = openai.Completion.create(
        model="text-davinci-003",
        prompt=prompt,
        temperature=0.1,
        max_tokens=10
    )
    response = response.choices[0].text.strip()
    if response not in options: response = response_aligner(options, response)
    return response

def yes_no(message):
    return assess_response(question="Yes or No?", response=message, options=["Yes", "No"], context="", uses_scale=False)

def questionnaire_chooser(message):
    options = available_questionnaires + ["Suicidal", "None of the previously mentioned options"]

    topic = assess_response(
        question = initial_message,
        response = message,
        options = options,
        uses_scale = False
        )
    
    match topic.lower():
        case "suicidal":
            state = ["initial"]
            response = f"""
                It sounds like you are in a crisis, or require immediate assistance. If this is the case, please call 911.
                If would like to know more, resources can be found here: {services}.
                """
        case "none of the previously mentioned options": 
            state = ["initial"]
            response = f"""
                At the moment I can only help with {listify(available_questionnaires)}.
                If you are in a crisis, or require immediate assistance, please call 911.
                If you are looking for help with another mental health topic, more resources can be found here: {services}.
                """
        case _:
            state = ["potential", topic.lower()]
            response = f"It sounds like you're having trouble with {topic.lower()}. Is that correct?" 
    
    return state, response

def get_question(questionnaire, number):
    return questionnaires.questions[questionnaire][number]

def empath_response(message, next_message):
    prompt = f"""
        You a therapist talking to a patient about their mental health issues.
        They have just told you {message}.
        You are about to say {next_message} to them.
        Before you do so, say something empathetic in one or two sentences."""
    response = openai.Completion.create(
        model="text-davinci-003",
        prompt=prompt,
        temperature=1,
        max_tokens=30
    )
    response = response.choices[0].text.strip().rstrip(".")
    return response

def generate_response(user, message, empath=True):
    
    user_id, state = user["user_id"], user["state"]

    match state[0]:    
        case "initial":
            state, response = questionnaire_chooser(message)
        
        case "potential":
            if yes_no(message) == "Yes": 
                state = ["screening", state[1], 1]
                response = get_question(state[1], 0)
            else:
                state, response = ["initial"], "I think I misunderstood you. Let's try again... can you please tell me about the problems you are experiencing?"

        case "screening":
            state, question = state, get_question(state[1], state[2])
            
            if question == questionnaires.END:
                state = ["finished"]
                response = f"I have finished my evaluation, and that is all I can do for you today. Please refresh my page if you would like to start again. If you would like to give your clinician access to my assessment, please give them reference number #{user_id}."
            else: 
                state[2] += 1
                response = question

        case "finished":
            state, response = state, ""

    if empath is True: response = f"{empath_response(message, next_message=response)}. " + response

    return state, response

if __name__ == '__main__':
    user = {"user_id": 0, "state": ["initial"]}
    while True:
        message = input("YOU: ")
        user["state"], response = generate_response(user, message)
        print("BRAINY: ", response)