import argparse
import json
import os
import random
import time
import requests
import openai


from levels.constants import all_objecsts, base_ingridients, capacity
from levels.utils import compute_dependency
from overcooked.game import World
from openai import OpenAI
from mistralai import Mistral


# Client
api_key = os.getenv("AVIOR_API_KEY") # Not support yet
api_base = "http://avior.mlfoundry.com/live-inference/v1"
client = OpenAI(
    api_key=api_key,
    base_url=api_base
)

# api_key = os.environ["MISTRAL_API_KEY"]
# client = Mistral(api_key=api_key)


key_id = 0
def rules(env, notes=True):
    prompt = ''
    prompt += 'The available actions are :\n'
    prompt += '1) goto: goto a tool location \n'
    prompt += '\t\tExample: goto_agent0_storage0 \n'
    prompt += '2) get: get some object from a tool \n'
    prompt += '\t\t Example: get_agent0_beef_storage0 \n'
    prompt += '3) put: put some object into a tool \n'
    prompt += '\t\t Example: put_agent0_blender0 \n'
    prompt += '4) activate: activate the tool to cook all ingredients inside the tool into a different tools \n'
    prompt += '\t\t Example: activate_agent0_blender0 \n'
    prompt += '5) noop: not performing any actions \n'
    prompt += '\t\t Example: noop_agent0 \n'

    prompt += 'Sometimes the system will give you error messages. Please consider these error messages when executing actions.  \n'
    prompt += 'You need to specify action for your agent. You can only specify one action at a time. \n'
    # prompt += 'You need to specify action for all of the agents, **except human**. They all have different agent numbers. Do not assign actions to the same agent more than once. \n'
    prompt += '\n'

    if notes:
        prompt += 'When the tools reach its capacity, you need to take stuff out. Otherwise, you cannot put items inside.\n'
        prompt += 'When you are holding objects, you cannot get any more objects. \n'
        prompt += "When you are holding objects, you cannot activate tools. \n"
        prompt += "Afer you cooked a required dish, you need to put it into the servingtable. \n"
        prompt += 'You can only pick up objects from the tool location, if you are located at the tool location. \n'
        prompt += "When you activate any tools, make sure all the items inside the tool are respecting the recipes. Otherwise, you will cook waste. Avoid waste at all cost. \n"

        prompt += "*** You should mix salad in the mixer. To make salad you should chop veggies first. *** \n"
        prompt += "*** If the tool is occupied, indicated by the occupy() predicate, you cannot get objects from it or put objects into it. *** \n"
        prompt += "*** The food orders are keep coming. You should finish as many dishes as possible and finish every dish as soon as possible. Please deliver the order to the serveringtable when it is finished. *** \n"
        prompt += "*** The dish will expire after the lifetime reaches 0 and it's not at the serveringtable. Please avoid this. *** "

    prompt += 'Here are the recipes: \n'
    prompt += recipes(env)

    prompt += generate_tool_descriptions(env)
    prompt += '\n\n'

    return prompt

def recipes(env: World):
    import copy
    tasks =  copy.deepcopy(env.task_manager._all_tasks)
    tasks.append('porkMeatcake')
    required_components = []

    for task in tasks:
        required_components.extend(compute_dependency(task)[0])
    required_components = set(required_components)

    prompt = '\n'
    with open('assets/recipe.json', 'r') as f:
        recipe = json.load(f)

    task_related_objects = set()
    task_related_tools = set()
    for dish, value in recipe.items():
        if dish in required_components:
            ingredients = value['ingredients']
            location = value['location']
            prompt += f'Cook {dish} at: \n'
            prompt += f' -- location: {location} \n'
            prompt += f' -- with ingredients: '
            task_related_objects.add(dish)
            task_related_tools.add(location)
            for ingredient in ingredients:
                task_related_objects.add(ingredient)
                prompt += f'    {ingredient}, '
            prompt += '\n'

    prompt += 'The following objects are available: \n'
    for idx, item in enumerate(task_related_objects):
        prompt += f' --{idx+1}) {item} \n'
    prompt += "The objecsts are cooked using tools or are just base ingredients. \n"

    prompt += "Among them, the following are base ingredients: \n"
    cnt = 1
    for idx, item in enumerate(task_related_objects):
        if item in base_ingridients:
            prompt += f" --{cnt}) {item} \n"
            cnt += 1
    prompt += "You can only obtain base ingredients from the storage initially.  \n"

    prompt += 'Additional rules: \n'
    for tool_name, tool in env.name_mapping.items():
        cap = capacity[tool_name[:-1]]
        num = cap
        if num == -1:
            num = 'infinite'

        prompt += f'You can place up to {num} item into the {tool_name} \n'
        prompt += f'You can place up to {num} item into the {tool_name} \n'

    return prompt

def generate_tool_descriptions(env: World):
    prompt = '** Only ** the following tools are available: \n'
    for tool_name, tool in env.name_mapping.items():
        prompt += f'{tool_name}, '

    prompt += 'You cannot pick up these tools. You can only use those tools at the corresponding location.'

    prompt += '\n'
    return prompt


def prepend_prompt(prompt, add, verbose=True):
    if verbose:
        print(add)
    return prompt + add

def prepend_history(history, add, role='user', verbose=True):
    if verbose:
        print(f'\n\n[[{role}]]\n\n' + add)
    assert role in ['user', 'assistant']
    history.append((role, add))
    return history


def query_llm(history, max_tokens=100, temperature=0.0, stop=['\n', '\n\n'], model=None):
    # print("HISTORY: ", history)
    if type(history) == str:
        history = [('user', history)]

    chat_history = []
    for i in history:
        if i[0] == 'user':
            chat_history.append({
                'role': 'user',
                'content': i[1]
            })
        elif i[0] == 'assistant':
            chat_history.append({
                'role': 'assistant',
                'content': i[1]
            })
        else:
            print(f"Role: {i[0]}")
            raise NotImplementedError

    # print(chat_history)

    # response = client.chat.complete(
    response = client.chat.completions.create(
                model=model,
                messages=chat_history,
                temperature=temperature,
                max_tokens=max_tokens,
                stop=stop
            )

    return response.choices[0].message.content


# if __name__ == '__main__':
#     args_parser = argparse.ArgumentParser(description='')
#     args_parser.add_argument('--max', type=int, default=1000)
#     args = args_parser.parse_args()

#     # Load your API key from an environment variable or secret management service
#     openai.api_key = open('./key.txt', 'r').read()


#     prompt = open('./assets/prompt_1.txt', 'r').read()
#     ret = query_llm(prompt + '\n Goal: make a meatcake with beef and flour',
#                     max_tokens=args.max)
#     print(ret)
