from agent import Agent
from utils.llm import rules, prepend_history, query_llm
from levels.utils import convert_to_prompt
from prompts.planner_agent import planner_info_prompt, planner_examples, executor_info_prompt, executor_examples

class PlannerAgent(Agent):
    def __init__(self, env, model, agent_id, total_num_agents, with_feedback, with_notes, look_ahead_steps, budget, costs):
        self.env = env
        self.model = model
        self.agent_id = agent_id
        self.total_num_agents = total_num_agents
        self.with_feedback = with_feedback
        self.with_notes = with_notes
        self.look_ahead_steps = look_ahead_steps

        self.budget = budget
        self.costs = costs
        self.history, self.feedback, self.suggestions = self.initialize_prompt()
        self.initial_history_length = len(self.history)
        self.prompt_history = []
        self.previous_actions = []

    def initialize_prompt(self):
        pre_prompt = ("user" , rules(self.env, self.with_notes))
        info_prompt = ("user", planner_info_prompt.format(total_num_agents=self.total_num_agents, budget=self.budget, costs=self.costs_to_txt(self.costs)))
        history = [pre_prompt] + [info_prompt] + planner_examples

        if self.with_feedback:
            feedback = '-execution error messages:\n  --  []\n'
            suggestions = '-execution suggestions:\n  --  []\n'
        else:
            feedback = ''
            suggestions = ''

        return history, feedback, suggestions

    def plan(self, obs, step, cost, budget, verbose=False):

        # update history
        if self.with_feedback and step != 0:
            self.feedback = '-execution error messages:\n  --  ' + str(self.env.feedback) + '\n'

        if self.with_notes and step !=0:
            self.suggestions = '-execution suggestions:\n  --  ' + str(self.env.suggestions) + '\n'
            if 'agent ids cannot be the same' in self.feedback:
                self.suggestions += f'  --  agent ids cannot be the same. \n'

        if step == 0:            
            prompt =  self.feedback + self.suggestions + convert_to_prompt(obs) + '-plan:\n'
        else:
            prompt = self.actions_to_txt(self.previous_actions) + self.feedback + self.suggestions + convert_to_prompt(obs) + f'-budget: ${cost}/{budget}\n' + '-plan:\n'
        print(f"STATE [step{step}]: ", convert_to_prompt(obs))
        
        # cap message length
        if len(self.history) < self.look_ahead_steps + self.initial_history_length:
            self.history = prepend_history(self.history, prompt, verbose=verbose)
        else:
            self.history = (self.history[:self.initial_history_length] + 
                          self.history[-(self.look_ahead_steps-1):])
            self.history = prepend_history(self.history, prompt, verbose=verbose)

        # GENERATE PLAN
        # print("PLANNER HISTORY:", self.history)
        plan = query_llm(self.history, max_tokens=200, model=self.model)

        self.history = prepend_history(self.history, plan, role='assistant', verbose=verbose)

        return plan


    

class ExecutorAgent(Agent):
    def __init__(self, env, model, agent_id, total_num_agents, with_feedback, with_notes, look_ahead_steps):
        super().__init__(env, model, agent_id, total_num_agents, with_feedback, with_notes, look_ahead_steps)

        self.history, self.feedback, self.suggestions = self.initialize_prompt()
        self.initial_history_length = len(self.history)
        self.prompt_history = []

    def initialize_prompt(self):
        pre_prompt = ("user" , rules(self.env, self.with_notes))
        info_prompt = ("user", executor_info_prompt.format(total_num_agents=self.total_num_agents, agent_id=self.agent_id))
        history = [pre_prompt] + executor_examples + [info_prompt]

        if self.with_feedback:
            feedback = '-execution error messages:\n  --  []\n'
            suggestions = '-execution suggestions:\n  --  []\n'
        else:
            feedback = ''
            suggestions = ''

        return history, feedback, suggestions
    
    def step(self, obs, step, plan, verbose=False):

        # update history
        if self.with_feedback and step != 0:
            self.feedback = '-execution error messages:\n  --  ' + str(self.env.feedback) + '\n'

        if self.with_notes and step !=0:
            self.suggestions = '-execution suggestions:\n  --  ' + str(self.env.suggestions) + '\n'
            if 'agent ids cannot be the same' in self.feedback:
                self.suggestions += f'  --  You can only control and plan the actions for agent{self.agent_id}. \n'
                    
        
        prompt =  self.feedback + self.suggestions + convert_to_prompt(obs) + '-plan:\n' + plan.strip() + '\n\n-action:\n'
        # cap message length
        if len(self.history) < self.look_ahead_steps + self.initial_history_length:
            self.history = prepend_history(self.history, prompt, verbose=verbose)
        else:
            self.history = (self.history[:self.initial_history_length] + 
                          self.history[-(self.look_ahead_steps-1):])
            self.history = prepend_history(self.history, prompt, verbose=verbose)

        # GENERATE ACTION
        # print(f"EXECUTOR HISTORY [agent{self.agent_id}]:", self.history)
        action = query_llm(self.history, model=self.model)
        # print(f"ACTION[agent{self.agent_id}]:", action)

        try:
            parsed_actions = self.extract_actions(action)
            if parsed_actions:
                parsed_actions = [parsed_actions[0]]
        except:
            parsed_actions = []

        if parsed_actions:
            self.update_history(parsed_actions, role='assistant', verbose=verbose)

        if parsed_actions:
            return [parsed_actions[0]]
        return parsed_actions