import json
import requests
import re

# config here
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:
        pairs = f.readlines()
    
    texts = []
    for pair in pairs:
        A = pair.split(', ')[0]
        B = pair.split(', ')[1]
        texts.append([A, B])
    return texts

def raw_query_gpt(prompt):
    data_any= {
        "model": "gpt-3.5-turbo",   
        "messages": [
            {"role": "system", "content": "I want you to act as an AI creative assistant helping users generate images using Midjourney. Your job is to generate a prompt for text-to-image generation model. You are good at text-image concept grounding, visual concept understanding, text-to-image diffusion models' prompt engineering. Don't be verbose in your answers, only provide the proposed prompt itself and avoid including explanations.\nI've provided you with a pair of visual concepts, and I want you to provide me with detailed prompts to feed into Midjourney that one given concept would show up with the other one in the corresponding generated image. It will be your job to come up with five prompts that closely integrate two concepts visually, yet each remains a distinct entity. Remember there is NO extra object or human in the prompt, NOR any explanatory sentences, according to the user's specifications."},
            {"role": "user", "content": "Now, the given two visual concepts are: Iced coke, teacup"},
            {"role": "assistant", "content": "Ice coke bubbling in a teacup.\nA teacup of ice coke.\nIce coke poured into a teacup\nIced coke in a teacup.\nTeacups filled with iced coke, cold fogging up."},
            {"role": "user", "content": "Now, the given two visual concepts are: Hainan Gibbon, Statue of Liberty"},
            {"role": "assistant", "content": "A tiny Hainan Gibbon perched atop the Statue of Liberty's torch.\nA Hainan Gibbon clinging to the Statue of Liberty's crown.\nThe Statue of Liberty adorned with a Hainan Gibbon.\nA Hainan Gibbon resting on the Statue of Liberty's shoulder.\nThe Hainan Gibbon swinging from the Statue of Liberty's tablet."},
            {"role": "user", "content": "Now, the given two visual concepts are: Panda, Fish scales"},
            {"role": "assistant", "content": "A panda covered in sleek fish scales.\nThe panda lies on the ground, its fish scales shimmering.\nA panda with black and white fish scales.\nA panda completely covered in fish scales.\nA panda with a body of fish scales."},
            {"role": "user", "content": "Now, the given two visual concepts are:   Lumberjack Shirt, Silk Robe"},
            {"role": "assistant", "content": "Tie the silk robe around the waist of the lumberjack shirt.\nDrape the silk robe over the lumberjack shirt.Drape the silk robe over the lumberjack shirt, exposing the shirt collar.\nSlip the silk robe over the lumberjack shirt.\nWearing the lumberjack shirt over the silk robe.\nLumberjack shirt leaking out of open collar of silk robe."},
            {"role": "user", "content": ""}
        ]
    }
    data_any["messages"][-1]["content"] = "Now, the given two visual concepts are: " + prompt
    try:
        resp = requests.post(url, headers=headers, data=json.dumps(data_any)).json()
    except:
        print("Error")
        return ""
    
    if resp.get('error'):
        print(resp['error'])
        return ""
    
    return resp['choices'][0]['message']['content']


def query_gpt(texts):
    dicts = []
    for i in range(len(texts)):
        print("Processing " + str(i) + "th pair")
        ans = raw_query_gpt(texts[i][0] + ", " + texts[i][1])
        while not ans:
            ans = raw_query_gpt(texts[i][0] + ", " + texts[i][1])
    
        dict = {}
        dict['A'] = texts[i][0]
        dict['B'] = texts[i][1]
        dict['prompt'] = ans.split('\n')
        dicts.append(dict)
        print(dict)
        print()
            
    with open(FILE_PATH[:-9] + '5sentences.json', 'w') as outfile:
        json.dump(dicts, outfile)      
            
if __name__ == '__main__':
    FILE_PATH = '' + '_raw_pairs.txt'
    print(FILE_PATH[:-9] + '5sentences.json')
    texts = process_data(FILE_PATH)
    query_gpt(texts)
    