import random
from levels.utils import convert_to_prompt

possible_actions = ['noop', 'goto', 'put', 'get', 'activate']

class RandomAgent():

    def __init__(self, env, total_num_agents, seed=42):
        random.seed(seed)
        self.total_num_agents = total_num_agents
        self.env = env

    def step(self, obs, verbose=False):
        available_actions = self.env.available_actions()
        # print(convert_to_prompt(obs))
        print(available_actions)

        actions = []
        for agent_id in range(self.total_num_agents):
            raw_action = random.choice(available_actions[agent_id])

            # format action
            action = self.format_action(raw_action, agent_id)
            actions.append(action)
        return actions
    
    def format_action(self, raw_action, agent_id):
        raw_action = raw_action.split(" ")
        if len(raw_action) == 1:
            return raw_action[0] + "_" + f"agent{agent_id}"
        elif len(raw_action) == 2:
            return raw_action[0] + "_" + f"agent{agent_id}" + "_" + raw_action[1]
        elif len(raw_action) == 3:
            return raw_action[0] + "_" + f"agent{agent_id}" + "_" + raw_action[1] + "_" + raw_action[2]
        else:
            raise ValueError("Invalid action")