import json
import requests
import re

# config here
FILE_PATH = "your path"
url = 'https://your_server/api/v1/assistant/any'
headers = {
    "Authorization": "Bearer " + "your_token",
    "Content-type": "application/json",
}


def process_data(file_path):
    dicts = []
    with open(file_path, 'r') as f:
        dicts = json.load(f)
    
    texts = []
    for dict in dicts:
        A = dict['A']
        B = dict['B']
        prompt = dict['prompt']
        texts.append([A, B, prompt])
    return texts

def gen_prompt(A, B):
    data_any= {
        "model": "gpt-3.5-turbo",   
        "messages": [
            {"role": "system", "content": "Task: Create a visual scene with a short sentence so that a given word fit in another."},
            {"role": "user", "content": "iced coke, teacup"},
            {"role": "assistant", "content": "Ice coke bubbling in a teacup."},
            {"role": "user", "content": "mammoth, spaceship"},
            {"role": "assistant", "content": "Mammoth standing in a spaceship."},
            {"role": "user", "content": "orange juice, teapot"},
            {"role": "assistant", "content": "An upside-down teapot pouring orange juice."},
            {"role": "user", "content": ""}
        ]
    }
    data_any["messages"][-1]["content"] = A + ", " + B
    resp = requests.post(url, headers=headers, data=json.dumps(data_any)).json()
    
    while resp.get('error'):
        print('error')
        resp = requests.post(url, headers=headers, data=json.dumps(data_any)).json()
    
    print("Prompt for ", A, " and ", B, ":")
    print(resp['choices'][0]['message']['content'])
    print()
    
    return resp['choices'][0]['message']['content']

def describe(A):
    data_any= {
        "model": "gpt-3.5-turbo",   
        "messages": [
            {"role": "system", "content": "Task: A concise description of what a given noun itself looks like, such as color, shape, or other visual features. As short as possible. No background. No other objects."},
            {"role": "user", "content": "coke"},
            {"role": "assistant", "content": "Dark brown liquid"},
            {"role": "user", "content": "tiger"},
            {"role": "assistant", "content": "Orange coat with black stripes, muscles, and claws."},
            {"role": "user", "content": "sushi"},
            {"role": "assistant", "content": "Rolled rice with colorful ingredients."},
            {"role": "user", "content": ""}
        ]
    }
    data_any["messages"][-1]["content"] = A
    try:
        resp = requests.post(url, headers=headers, data=json.dumps(data_any)).json()
    except:
        print("Error")
        return ""
    
    if resp.get('error'):
        print('error')
        return ""
    
    print("Description of ", A, ":")
    print(resp['choices'][0]['message']['content'])
    
    return resp['choices'][0]['message']['content']

# return a dictionary of texts {{"A":"", "B":"", "Prob":""}, {}, ...}
def sol(texts):
    dicts = {}
    for i in range(len(texts)):
        dict = {}
        print("Processing " + str(i) + "th pair")
        
        prompt = texts[i][2]
        dict["prompt"] = prompt
        
        dict["first_prompt"] = texts[i][1] # The container or the background
        dict["first_describe"] = describe(texts[i][1])
        while dict["first_describe"] == "":
            dict["first_describe"] = describe(texts[i][1])
        dict["second_prompt"] = texts[i][0] # The beverage or the prospect
        dict["second_describe"] = describe(texts[i][0])
        while dict["second_describe"] == "":
            dict["second_describe"] = describe(texts[i][0])
    
        dicts["prompt" + str(i)] = dict
            
    with open(FILE_PATH[:-5] + '-firstOrSecond_v3.json', 'w') as outfile:
        json.dump(dicts, outfile)      
            
if __name__ == '__main__':
    texts = process_data(FILE_PATH)
    sol(texts)
    