from my_agent.agent import Agent
from my_agent.agent_debate import DebateAgent
from my_agent.agent_orchestrator import OrchestratorAgent
from my_agent.agent_planner import PlannerAgent, ExecutorAgent
from my_agent.agent_random import RandomAgent
from my_agent.agent_auction import AuctionAgent, AuctionSystem
import copy

class AgentManager: 
    def __init__(self, env, model, total_num_agents, with_feedback, with_notes, look_ahead_steps, agent_type, budget, num_debate_rounds=None, costs=None):
        
        self.agents = []
        self.total_num_agents = total_num_agents
        self.agent_type = agent_type
        self.num_debate_rounds = num_debate_rounds

        self.budget = budget
        self.costs = costs
        self.cost = 0
        self.predictions = []

        if agent_type == 'orchestrator':
            self.orchestrator = OrchestratorAgent(env, model, -1, total_num_agents, with_feedback, with_notes, look_ahead_steps, budget, self.costs)
        elif agent_type == 'planner':
            self.planner = PlannerAgent(env, model, -1, total_num_agents, with_feedback, with_notes, look_ahead_steps, budget, self.costs)
        elif agent_type == 'random':
            self.random = RandomAgent(env, total_num_agents)
        elif agent_type == 'auction':
            # Create auction agents
            for agent_id in range(total_num_agents):
                agent = AuctionAgent(env, model, agent_id, total_num_agents, with_feedback, with_notes, look_ahead_steps)
                self.agents.append(agent)
            # Create the auction system
            self.auction_system = AuctionSystem(env, model, self.agents, self.costs)
        else:
            pass

        # Create agents for non-auction agent types
        if agent_type != 'auction':
            for agent_id in range(total_num_agents):
                if agent_type == 'individual':
                    agent = Agent(env, model, agent_id, total_num_agents, with_feedback, with_notes, look_ahead_steps)
                elif agent_type == 'debate':
                    agent = DebateAgent(env, model, agent_id, total_num_agents, with_feedback, with_notes, look_ahead_steps, num_debate_rounds)
                elif agent_type == 'orchestrator':
                    break
                elif agent_type == 'planner':
                    agent = ExecutorAgent(env, model, agent_id, total_num_agents, with_feedback, with_notes, look_ahead_steps)
                elif agent_type == 'random':
                    break
                else: 
                    raise ValueError("Invalid agent type")
                self.agents.append(agent)
            
            
    def step(self, obs, step, verbose=False): 

        if self.agent_type == 'individual':
            plan = self.individual_step(obs, step, verbose=verbose)
        elif self.agent_type == 'debate':
            plan = self.debate_step(obs, step, verbose=verbose)
        elif self.agent_type == 'orchestrator':
            plan = self.orchestrate_step(obs, step, verbose=verbose)
        elif self.agent_type == 'planner':
            plan = self.planner_step(obs, step, self.cost, self.budget, verbose=verbose)
        elif self.agent_type == 'random':
            plan = self.random_step(obs,verbose=verbose)
        elif self.agent_type == 'auction':
            plan = self.auction_step(obs, step, verbose=verbose)
        else:
            raise ValueError("Invalid agent type")

        self.cost += self.compute_cost(plan)
        print("COST: ", self.cost)
        return plan

    def individual_step(self, obs, step, verbose=False): 
        """
        Generate a plan for each agent individually for a given step.

        Args:
            obs (object): The current observation of the environment.
            step (int): The current step number in the simulation.
            verbose (bool, optional): If True, enables verbose logging. Defaults to False.

        Returns:
            list: A list of actions for all agents combined.
        """
        plan = []
        for agent_id in range(self.total_num_agents):
            agent = self.agents[agent_id]
            actions = agent.step(obs, step, verbose=verbose)
            plan.extend(actions)
        return plan
    
    def debate_step(self, obs, step, verbose=False):
        
        # First round - all agents take initial actions
        actions = []
        for agent_id in range(self.total_num_agents):
            agent = self.agents[agent_id]
            action = agent.step(obs, step, verbose=verbose)  # Use step() instead of debate() for initial actions
            actions.extend(action)

        # Subsequent rounds - agents debate about their actions
        debate = actions
        for round in range(self.num_debate_rounds):
            current_debate = []
            for agent_id in range(self.total_num_agents):
                agent = self.agents[agent_id]
                action = agent.debate(obs, step, round, debate, verbose=verbose)  # Pass previous round's actions
                current_debate.extend(action)
            debate = current_debate  # Update debate with current round's actions
            print(f"ACTIONS at round {round}:", debate) 

        return debate

    def orchestrate_step(self, obs, step, verbose=False):
        """
        Generate a plan for the orchestrator agent for a given step.

        Args:
            obs (object): The current observation of the environment.
            step (int): The current step number in the simulation.
            verbose (bool, optional): If True, enables verbose logging. Defaults to False.

        Returns:
            string: A plan for each of the executor agents.
        """
        # PLAN  ACTIONS
        plan = self.orchestrator.step(obs, step, self.cost, self.budget, verbose=verbose)
        
        # Create a deep copy of the success rates before appending
        self.predictions.append(copy.deepcopy(self.orchestrator.agent_action_success_rates))
        
        return plan
    
    def planner_step(self, obs, step, cost, budget, verbose=False):
        """
        Generate a plan for the planner agent for a given step.
        """

        plan = self.planner.plan(obs, step, cost, budget, verbose=verbose)
        print("PLANNER PLAN:", plan)
        
        actions = []
        for agent_id in range(self.total_num_agents):
            print(f"Querying agent{agent_id}")
            agent = self.agents[agent_id]
            action = agent.step(obs, step, plan, verbose=verbose)
            actions.extend(action)
            print(f"Agent{agent_id}: ", action)

        self.planner.previous_actions = actions

        return actions  
    
    def random_step(self, obs, verbose=False):

        actions = self.random.step(obs, verbose=verbose)
        return actions
    
    def auction_step(self, obs, step, verbose=False):
        """
        Implements an auction-based approach where agents bid for tasks using a virtual currency system.
        
        Uses the AuctionSystem class to handle the auction process.
        
        Args:
            obs (dict): The current observation of the environment.
            step (int): The current step number in the simulation.
            verbose (bool, optional): If True, enables verbose logging. Defaults to False.
            
        Returns:
            list: A list of actions for all agents based on auction results.
        """
        return self.auction_system.run_auction(obs, step, verbose)
    
    def compute_cost(self, actions):
        step_cost = 0
        for i in range(len(actions)):
            full_action = actions[i]
            action = full_action.split('_')[0]
            if action == "noop":
                continue
            else:
                step_cost += self.costs[i][action]
        return step_cost





