from utils.llm import rules
from my_agent.agent import Agent
from prompts.debate_agent_action_only import debate_info_prompt, debate_examples
from utils.llm import rules, prepend_history, query_llm
from levels.utils import convert_to_prompt

class DebateAgent(Agent):
    def __init__(self, env, model, agent_id, total_num_agents, with_feedback, with_notes, look_ahead_steps, num_debate_rounds):
        # Initialize parent class first
        self.num_debate_rounds = num_debate_rounds
        super().__init__(env, model, agent_id, total_num_agents, with_feedback, with_notes, look_ahead_steps)
        

    def initialize_prompt(self):
        pre_prompt = ("user", rules(self.env, self.with_notes))
        print(">>> Number of debate rounds:", self.num_debate_rounds)
        info_prompt = ("user", debate_info_prompt.format(
            total_num_agents=self.total_num_agents, 
            agent_id=self.agent_id,
            debate_rounds=self.num_debate_rounds
            ))

        # history = [pre_prompt] + debate_examples + [info_prompt]
        history = [pre_prompt] + debate_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, 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) + '-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
        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


    
    def debate(self, obs, step, round, previous_actions, verbose=False):
        """
        Executes a single debate round for the agent based on the given observation.

        Args:
            obs (dict): The observation data for the current step.
            step (int): The current step number.
            round (int): The current debate round number.
            verbose (bool): Whether to print verbose output.

        Returns:
            list: A list of actions to take.
        """

        previous_actions_text = ""
        for agent_id, action in enumerate(previous_actions):
            if action:  # Only add if there are actions
                previous_actions_text += f"Agent {agent_id}'s action: {action}\n"

        prompt = f"""-Debate round {round+1}:
        {previous_actions_text}

        -your updated action:
        """
        
        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)

        # print("DEBATE PROMPT: ", self.history)
        action = query_llm(self.history, model=self.model)

        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

      

    
