import re 

from utils.llm import rules, prepend_history, query_llm
from levels.utils import convert_to_prompt
from prompts.individual_agent import individual_examples, individual_info_prompt


class Agent():
    def __init__(self, env, model, agent_id, total_num_agents, with_feedback, with_notes, look_ahead_steps): 
        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.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", individual_info_prompt.format(total_num_agents=self.total_num_agents, agent_id=self.agent_id))
        examples = [(e[0], e[1].replace("agent0", f"agent{self.agent_id}")) for e in individual_examples]
        history = [pre_prompt] + 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, agent_tasks=None): 
        """
        Executes a single step for the agent based on the given observation.

        Args:
            obs (dict): The observation data for the current step.
            step (int): The current step number.
            verbose (bool, optional): If True, prints detailed debug information. Defaults to False.
            agent_tasks (list, optional): List of tasks assigned to this agent through auction. Defaults to None.

        Returns:
            list: A list of parsed actions generated by the agent.
        """

        # 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'
        
        # Add auction-assigned tasks to the prompt if provided
        auction_tasks_prompt = ""
        if agent_tasks:
            auction_tasks_prompt = f"-assigned tasks:\n"
            for task_id in agent_tasks:
                task_parts = task_id.split('_')
                parent_task = task_parts[0]
                action_type = task_parts[1]
                item = '_'.join(task_parts[2:]) if len(task_parts) > 2 else ""
                auction_tasks_prompt += f"  -- {action_type} {item} for {parent_task}\n"
                    
        prompt = self.feedback + self.suggestions + convert_to_prompt(obs) + auction_tasks_prompt + '-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"HISTORY[agent{self.agent_id}]:", len(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

    
    def update_history(self, actions, role='assistant', verbose=False):

        to_add = '\n'.join(actions)
        self.history = prepend_history(self.history, to_add, role, verbose=verbose)


    def extract_actions(self, text):
        # List of action types
        action_types = ["noop", "goto", "put", "activate", "get"]
        
        # Pattern for the actions
        pattern = r'((' + '|'.join(action_types) + r')_agent\d+(_[a-zA-Z0-9_]+)?)'
        
        matches = re.findall(pattern, text)
        
        # Extracting just the full action names from the returned tuples
        actions = [match[0] for match in matches]
        return actions
    

    def actions_to_txt(self, actions):
        txt = ""
        txt += f"-previous actions:\n"
        for i, action in enumerate(actions):
            txt += f"Agent{i}: {action}\n"
        return txt
    
    def costs_to_txt(self, costs):
        txt = ""
        # Get list of all actions
        actions = list(costs[0].keys())
        
        # Calculate column widths
        action_width = max(len(action) for action in actions)
        cost_width = 8  # Width for cost values
        
        # Create header
        txt += "│ " + "Agent".ljust(5) + " │"
        for action in actions:
            txt += f" {action.ljust(action_width)} │"
        txt += "\n"
        
        # Add separator line
        txt += "├" + "─" * 7 + "┼"
        for _ in actions:
            txt += "─" * (action_width + 2) + "┼"
        txt = txt[:-1] + "┤\n"
        
        # Add data rows
        for agent_id, agent_costs in enumerate(costs):
            txt += f"│ {str(agent_id).ljust(5)} │"
            for action in actions:
                cost = f"${agent_costs[action]:.2f}"
                txt += f" {cost.ljust(action_width)} │"
            txt += "\n"
            
        # Add bottom border
        txt += "└" + "─" * 7 + "┴"
        for _ in actions:
            txt += "─" * (action_width + 2) + "┴"
        txt = txt[:-1] + "┘\n"
        
        return txt


