import copy

from my_agent.agent import Agent
from my_agent.agent_orchestrator import OrchestratorAgent
from my_agent.agent_planner import ExecutorAgent, PlannerAgent


class AgentManager: 
    def __init__(self, env, model, total_num_agents, with_feedback, with_notes, look_ahead_steps, agent_type, executor_models=None):
        
        self.agents = []
        self.total_num_agents = total_num_agents
        self.agent_type = agent_type
        self.price = 0
        self.current_plan = ""
        self.individual_costs = {}  # Track costs per agent

        self.predictions = []

        if agent_type == 'orchestrator':
            self.orchestrator = OrchestratorAgent(env, model, -1, total_num_agents, with_feedback, with_notes, look_ahead_steps)
            self.individual_costs['orchestrator'] = 0
        elif agent_type == 'planner':
            self.planner = PlannerAgent(env, model, -1, total_num_agents, with_feedback, with_notes, look_ahead_steps)
            self.individual_costs['planner'] = 0
        else:
            pass

        # Create agents for non-auction agent types
        for agent_id in range(total_num_agents):
            if agent_type == 'orchestrator':
                break
            elif agent_type == 'planner':
                # Split the executor models string into a list if it's a string
                if isinstance(executor_models, str):
                    executor_models_list = executor_models.split(',')
                else:
                    executor_models_list = executor_models
                
                # Use the corresponding executor model for this agent
                executor_model = executor_models_list[agent_id]
                agent = ExecutorAgent(env, executor_model, agent_id, total_num_agents, with_feedback, with_notes, look_ahead_steps)
                self.individual_costs[f'executor_{agent_id}'] = 0
            else: 
                raise ValueError("Invalid agent type")
            self.agents.append(agent)
            
            
    def step(self, obs, step, verbose=False): 

        if self.agent_type == 'orchestrator':
            plan, price = self.orchestrate_step(obs, step, verbose=verbose)
            self.individual_costs['orchestrator'] += price
        elif self.agent_type == 'planner':
            plan, price = self.planner_step(obs, step, verbose=verbose)
            self.individual_costs['planner'] += price
        else:
            raise ValueError("Invalid agent type")

        self.price += price

        return plan
    

    def orchestrate_step(self, obs, step, verbose=False):
        """
        The orchestrator agent will generate a plan of actions
        """

        # Get the actions from the orchestrator agent
        actions, price = self.orchestrator.step(obs, step, verbose=verbose)
        return actions, price
    
    def planner_step(self, obs, step, verbose=False):
        """
        Generate a general plan with the planner agent for all the agents when a new dish is added, removed, or completed. 
        Then generate the actions with the executor agents.
        """
        total_price = 0  # Initialize price

        # Check if there are dish-related events that require planning
        should_plan = (step == 0 or  # First step always needs a plan
                      obs.just_new_task or  # New dish arrived
                      obs.just_failed or  # Dish removed
                      obs.task_just_success)  # Dish completed
        
        # Generate a new plan if needed
        if should_plan:
            if verbose:
                print(f"Step {step}: Generating new plan due to dish state change")
            # Use the planner's plan method rather than step
            self.current_plan, price = self.planner.plan(obs, step, verbose=verbose)
            print(f"PLAN[step={step}]: ", self.current_plan)
            total_price += price

        # Get actions from all executor agents based on current plan
        all_actions = []
        
        for agent_id, agent in enumerate(self.agents):
            # Each executor gets the current plan and decides its own action
            agent_actions, price = agent.step(obs, step, self.current_plan, verbose=verbose)
            if agent_actions:
                all_actions.extend(agent_actions)
            total_price += price
            self.individual_costs[f'executor_{agent_id}'] += price
        
        return all_actions, total_price





