import os,openai,json,time

def check_api_key_validity(api_key):
   try:
        openai.api_key = api_key
        print("OpenAI API key is valid")
   except openai.OpenAIError as e:
        print("Invalid OpenAI API key")
        exit()


# Much cheaper and powerful than Davinci
def call_AI_chatGPT(prompt):
        model = "gpt-3.5-turbo"
      # generate response
        generated=False
        while(not generated):
            try:
                time.sleep(2)
                response = openai.ChatCompletion.create(
                    model=model,
                    messages =[
                        {'role':'system','content':'You are an Argentine assistant. When given an instruction, answer the instruction in spanish and write the response of the instruction and nothing else. Use Argentine slang like boludo, che, pelotudo, gil, laburo, forro, sacarla barata, mala leche, al pedo, gil de cuarta, quilombo, estas en pedo, ya fue, etc. Do not always start with che or boludo.'},
                        {'role':'user','content':prompt}
                        ],
                    temperature=1.0,
                    max_tokens=1024,request_timeout=40)
                generated=True
            except Exception as e:
                print("%s: Trying again..." % str(e))
                pass
        return response.choices[0]['message']['content']

# import api key
api_key = os.environ.get("OPENAI_API_KEY")
if (api_key is None) or (len(api_key)==0): # try to load apikey from file
    try:
        api_key=open('api-key.txt','rb').read().strip().decode()
    except:
        print("Couldn't load OpenAI Api key, please load it in OPENAI_API_KEY env variable, or alternatively in 'api-key.txt' file.")
        exit(0)
else: print('Loaded api key from environment variable.')

check_api_key_validity(api_key)

# Original ALPACA prompt
PROMPT_DICT = {
    "prompt_input": (
        "You are an Argentine assistant. When given an instruction, answer the instruction in spanish and write the response of the instruction and nothing else. Use Argentine slang like boludo, che, pelotudo, gil, laburo, forro, mala leche, al pedo, gil de cuarta, quilombo, estas en pedo, ya fue. Do not always start with Che or boludo. Answer in a colloquial manner. Below is an instruction that describes the task, paired with an input that provides further context, and the output response. "
        "Rewrite the response that appropriately completes the request, taking in acount that you are Argentine, and respond as such.\n\n"
        "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response: \n{output}\n\n### Rewrite: \n"
    ),
    "prompt_no_input": (
        "You are an Argentine assistant. When given an instruction, answer the instruction in spanish and write te response to the instruction and nothing else. Use Argentine slang like boludo, che, pelotudo, gil, laburo, forro, mala leche, al pedo, gil de cuarta, quilombo, estas en pedo, ya fue, etc. Do not overuse slang. Below is an instruction that describes a task, and the output response. "
        "Rewrite the response that appropriately completes the request, taking in acount that you are Argentine, and respond as such, but do not always start with che or boludo.\n\n"
        "### Instruction:\n{instruction}\n\n### Input:\n{input}\n\n### Response: \n{output}\n\n### Rewrite: \n"
    ),
}

originalDataFile="alpaca_data_cleaned_spanish.json"

f=open(originalDataFile,"r")
data=json.load(f)
f.close()
print("Loaded %s data file with %d entries." % (originalDataFile,len(data)))

prompt_input, prompt_no_input = PROMPT_DICT["prompt_input"], PROMPT_DICT["prompt_no_input"]
index=0
for i in data:
    if len(i['input'])==0:
        prompt=prompt_no_input.format_map(i)
    else:
        prompt=prompt_input.format_map(i)
    i["output"]=call_AI_chatGPT(prompt).strip()
    js = json.dumps(i)
    print("Index: %d/%d: %s" % (index,len(data),js))
    index+=1
    e=open("dataset-alpaca-Arg.json","a")
    e.write("%s,\n" % js)
    e.close()