import json

import openai

functions = [
    {
        "name": "with_context",
        "description": (
            "Writes the question in a more verbose way by filling in missing data from the conversation, to make it more clear"
            ", imagine you are a third person that just got in on the conversation and heard only the last question. Make it easier to understand what it is about."
            "Replace words like 'there' with the actual location, 'that' with the actual subject, 'it' with the actual object, etc."
        ),
        "parameters": {
            "type": "object",
            "properties": {
                "contextualized_question": {
                    "type": "string",
                    "description": "The contextualized question",
                }
            },
            "required": ["contextualized_question"],
        },
    }
]


def contextualize_question(user_query: str, messages: list[str]):
    prompt = (
        f"The user said: {user_query}\n\n"
        "If the user is asking a question, make sure you contextualize it using the replies exchanged before, add relevant details to the question.\n"
        "The previous messages are:\n"
    )

    messages_for_context = messages[1:]
    for message in messages_for_context:
        prompt += f"<<{message['role']}>>{message['content']}<<{message['role']}>>\n"

    res = openai.ChatCompletion.create(
        model="gpt-3.5-turbo",
        messages=[
            {
                "role": "system",
                "content": "You are a helpful assistant adds context to vague questions.",
            },
            {"role": "user", "content": prompt},
        ],
        functions=functions,
        function_call={"name": "with_context"},  # force the function to be called
        temperature=0.1,
    )

    try:
        arguments = res["choices"][0]["message"]["function_call"]["arguments"]

        result_data = json.loads(arguments)

        return result_data["contextualized_question"]
    except Exception as error:
        print(error)
        return user_query