import requests
import argparse
from nltk import sent_tokenize
from .contexts import ManualContexts, AwesomePrompts, CustomPrompt

'''
This function wraps user's petition question with the adequate context to better orient the response of the language model
'''

class Wrapper:

    def wrapper(self, query:str, key: str, context:str, model:str="text-davinci-003") -> str:
        contexts = AwesomePrompts.dataset['act']
        contexts.extend(["doctor", "lawyer", "programmer"])
        contexts.extend(CustomPrompt.available_prompts)
        if not any(context==x for x in contexts):
            raise Exception("Sorry, that context is not implemented yet")

        url = 'https://api.openai.com/v1/completions'
        headers = {"Content-Type": "application/json; charset=utf-8", "Authorization":"Bearer "+key}
        if context=="doctor":
            context = ManualContexts.DoctorAdvice
        elif context=="lawyer":
            context = ManualContexts.Lawyer
        elif context=="programmer":
            context = ManualContexts.Programmer
        elif context.startswith('custom-'):
            custom_name = context[len('custom-'):]
            context = CustomPrompt(custom_name).prompt
        else:
            context = AwesomePrompts.dataset.filter(lambda x: x['act']==context)['prompt'][0]
            context = ' '.join(sent_tokenize(context)[:-1])

        myobj = {'model': model, 'prompt': context + " \"" + query+ "\"", "max_tokens":256, "temperature":0} # temperature is set to 0 by default since we want the most deterministic as possible responses
                                                                                                                    # max_tokens = 256 because we want a concrete explanation, better if it yes or no

        x = requests.post(url, headers =headers, json = myobj)
        
        return x.json()['choices'][0]['text']
