import base64
import requests
from pprint import pprint

api_key = "..."

headers = {
    "Content-Type": "application/json",
    "Authorization": f"Bearer {api_key}"
}

def prompt(command, situation):
    return \
        "You are a robot named iCubSim. " +\
        "You can move with two hands, you can touch and hit. " +\
        "You can perform the following actions: " +\
        "Smile(True or False), Sleep(miliseconds), Say(text), Hit(), StandardPosition() and WaveWithHand(). " +\
        "Please answer with json code. Return action parameters as a list. " +\
        "Please always use the Say action at least one time. " +\
        "Situation: There is a table in from of you. "+ situation + \
        " '" + command + "' " +\
        "Actions:"

def get(dictionary, attributes, default):
    for attribute in attributes:
        if attribute in dictionary:
            return dictionary[attribute]
    return default

def askGPT(
    command = "Remove the object from the table!",
    situation = "There is a ball on the table."
):
    global request
    request = prompt(command, situation)
    pprint(request)
    print()
    payload = {
        "model": "gpt-4-0613",
        "messages": [ {
            "role": "user",
            "content": request
        } ],
        "max_tokens": 300
    }
    global response
    response = requests.post(
        "https://api.openai.com/v1/chat/completions", 
        headers=headers, json=payload
    )
    pprint(response.json())
    print()
    global content
    try:
        content = response.json().get('choices')[0]['message']['content']
        true = True
        false = False
        allactions = eval(content)
        if isinstance(allactions,dict):
            allactions = get(allactions,["Actions","actions",'ActionSequence'],[value for value in allactions.values()])
        known = ['StandardPosition', 'Hit', 'Smile', 'Say', 'WaveWithHand']
        myactions = []
        for action in allactions:
            if get(action,['action','Action'],'') in known:
                myactions.append({
                    'action': get(action,['action','Action'],''),
                    'parameters': get(action,['parameters','Parameters'],[])
                })
        return myactions
    except:
        return [{'action' : 'Say', 'parameters' : ['I have no idea']}]

if __name__ == "__main__":  
    actions = askGPT(
        "Remove the object from the table!",
        "There is no object on the table." 
    )
    for action in actions:
        print(action['action']+"(",end="")
        for i, parameter in enumerate(action['parameters']):
            if i > 0:
                print(",",end="")
            print(parameter,end="")
        print(")")
            
