import collections
from dataclasses import dataclass, field
from typing import List, Dict, Set, Optional
import uuid

@dataclass
class Task:
    id: str = field(default_factory=lambda: str(uuid.uuid4()))
    type: str
    payload: Dict
    dependencies: List[str] = field(default_factory=list)
    priority: int = 0  # Higher value means higher priority

@dataclass
class Agent:
    id: str
    capabilities: Set[str]  # e.g., {'DATA_PROCESSING', 'MODEL_TRAINING'}
    can_handle_dependencies: bool = False
    current_load: int = 0  # Number of tasks currently assigned
    # In a real system, this might be a more complex metric like CPU/memory usage

class TaskRouter:
    """
    The TaskRouter is responsible for distributing tasks to optimal agents
    within Queen AIVA's evolving swarm. It considers task type, agent capabilities,
    load balancing, and task dependencies to ensure efficient processing.
    """
    def __init__(self):
        self._agents: Dict[str, Agent] = {}
        self._pending_tasks: collections.deque[Task] = collections.deque() # Tasks awaiting dependencies or agents
        self._completed_tasks: Set[str] = set() # To track completed task IDs for dependency resolution
        print("TaskRouter initialized for Queen AIVA's swarm. Ready to distribute.")

    def register_agent(self, agent: Agent):
        """
        Registers an agent with the router, making it available for task assignment.
        """
        if agent.id in self._agents:
            print(f"Warning: Agent {agent.id} already registered. Updating capabilities.")
        self._agents[agent.id] = agent
        print(f"Agent {agent.id} registered with capabilities: {agent.capabilities}, dependency handling: {agent.can_handle_dependencies}")

    def unregister_agent(self, agent_id: str):
        """
        Removes an agent from the router's pool.
        """
        if agent_id in self._agents:
            del self._agents[agent_id]
            print(f"Agent {agent_id} unregistered.")
        else:
            print(f"Warning: Agent {agent_id} not found for unregistration.")

    def _find_suitable_agents(self, task: Task) -> List[Agent]:
        """
        Identifies agents capable of handling the given task type.
        Prioritizes agents that can handle dependencies if the task has them.
        """
        suitable_by_type = [
            agent for agent in self._agents.values()
            if task.type in agent.capabilities
        ]

        if task.dependencies:
            # If task has dependencies, prioritize agents that can manage them
            dependency_capable_agents = [
                agent for agent in suitable_by_type
                if agent.can_handle_dependencies
            ]
            if dependency_capable_agents:
                return dependency_capable_agents
            # Fallback if no specific dependency-capable agent for this type, but they *can* handle the type

        return suitable_by_type

    def _select_agent(self, suitable_agents: List[Agent]) -> Optional[Agent]:
        """
        Selects the optimal agent from a list of suitable agents based on current load.
        """
        if not suitable_agents:
            return None

        # Sort by current_load to pick the least loaded agent
        # In a more advanced system, this could involve more complex metrics or prediction
        sorted_agents = sorted(suitable_agents, key=lambda agent: agent.current_load)
        return sorted_agents[0]

    def _check_dependencies_met(self, task: Task) -> bool:
        """
        Checks if all dependencies for a given task have been marked as completed.
        """
        if not task.dependencies:
            return True # No dependencies, so they are 'met'
        
        # All dependencies must be in the _completed_tasks set
        return all(dep_id in self._completed_tasks for dep_id in task.dependencies)

    def route_task(self, task: Task) -> Optional[str]:
        """
        Receives a task and attempts to route it to an optimal agent.
        Returns the ID of the assigned agent, or None if no agent could be found
        or if dependencies are not yet met.
        """
        if not self._check_dependencies_met(task):
            print(f"Task {task.id} (type: {task.type}) has unmet dependencies. Adding to pending queue.")
            self._pending_tasks.append(task)
            return None

        suitable_agents = self._find_suitable_agents(task)
        selected_agent = self._select_agent(suitable_agents)

        if selected_agent:
            selected_agent.current_load += 1
            print(f"Task {task.id} (type: {task.type}) routed to Agent {selected_agent.id}. Load: {selected_agent.current_load}")
            # In a real system, the task would be sent to the agent via a message queue or direct call
            return selected_agent.id
        else:
            print(f"No suitable agent found for Task {task.id} (type: {task.type}). Adding to pending queue.")
            self._pending_tasks.append(task)
            return None

    def notify_task_completion(self, task_id: str, agent_id: str):
        """
        An agent notifies the router that a task has been completed.
        This updates agent load and allows for dependent tasks to be routed.
        """
        if agent_id in self._agents:
            self._agents[agent_id].current_load = max(0, self._agents[agent_id].current_load - 1)
            print(f"Agent {agent_id} completed task {task_id}. New load: {self._agents[agent_id].current_load}")
        else:
            print(f"Warning: Agent {agent_id} not found for task completion notification.")
        
        self._completed_tasks.add(task_id)
        print(f"Task {task_id} marked as completed. Attempting to clear pending tasks.")
        self._process_pending_tasks()

    def _process_pending_tasks(self):
        """
        Attempts to route tasks from the pending queue whose dependencies are now met.
        """
        tasks_to_requeue = collections.deque()
        while self._pending_tasks:
            task = self._pending_tasks.popleft()
            if self._check_dependencies_met(task):
                print(f"Attempting to route previously pending task {task.id}.")
                routed_agent_id = self.route_task(task)
                if routed_agent_id is None: # Still couldn't route (e.g., no suitable agent available)
                    tasks_to_requeue.append(task)
            else:
                tasks_to_requeue.append(task) # Dependencies still not met, put back
        self._pending_tasks = tasks_to_requeue
        if self._pending_tasks: 
            print(f"{len(self._pending_tasks)} tasks remain in pending queue.")

# Example Usage for verification:
if __name__ == "__main__":
    router = TaskRouter()

    # 1. Register Agents
    agent_a = Agent(id="Agent_A", capabilities={"DATA_PROCESSING", "ANALYTICS"}, can_handle_dependencies=True)
    agent_b = Agent(id="Agent_B", capabilities={"DATA_PROCESSING"})
    agent_c = Agent(id="Agent_C", capabilities={"MODEL_TRAINING", "ANALYTICS"})
    agent_d = Agent(id="Agent_D", capabilities={"MODEL_TRAINING"}, can_handle_dependencies=True)

    router.register_agent(agent_a)
    router.register_agent(agent_b)
    router.register_agent(agent_c)
    router.register_agent(agent_d)

    print("\n--- Routing Tasks ---")

    # 2. Route Tasks - Test Task Type Routing & Load Balancing
    task1 = Task(type="DATA_PROCESSING", payload={"data": "raw_data_1"})
    task2 = Task(type="DATA_PROCESSING", payload={"data": "raw_data_2"})
    task3 = Task(type="MODEL_TRAINING", payload={"model_id": "M1"})
    task4 = Task(type="ANALYTICS", payload={"report_id": "R1"})
    task5 = Task(type="DATA_PROCESSING", payload={"data": "raw_data_3"})

    router.route_task(task1) # Should go to A (load 0) or B (load 0), let's say A
    router.route_task(task2) # Should go to B (load 0)
    router.route_task(task3) # Should go to C (load 0) or D (load 0), let's say C
    router.route_task(task4) # Should go to A (load 1) or C (load 1), let's say A (still lower if A was 1 and C was 1)
    router.route_task(task5) # Should go to A (load 1) or B (load 1), let's say B

    print(f"\nAgent A Load: {agent_a.current_load}")
    print(f"Agent B Load: {agent_b.current_load}")
    print(f"Agent C Load: {agent_c.current_load}")
    print(f"Agent D Load: {agent_d.current_load}")

    print("\n--- Routing Tasks with Dependencies ---")
    # 3. Test Dependencies
    task_dep_1 = Task(id="DEP_TASK_1", type="DATA_PROCESSING", payload={"data": "dep_data_1"})
    task_dep_2 = Task(id="DEP_TASK_2", type="ANALYTICS", payload={"analysis": "dep_analysis_2"}, dependencies=["DEP_TASK_1"])
    task_dep_3 = Task(id="DEP_TASK_3", type="MODEL_TRAINING", payload={"model": "dep_model_3"}, dependencies=["DEP_TASK_1", "DEP_TASK_2"])

    # DEP_TASK_2 and DEP_TASK_3 should be pending initially
    router.route_task(task_dep_2) # Dependencies unmet, goes to pending
    router.route_task(task_dep_3) # Dependencies unmet, goes to pending

    # Route DEP_TASK_1, which has no dependencies
    router.route_task(task_dep_1) # Should go to A or B

    print("\n--- Notifying Completion to Trigger Dependencies ---")
    # Notify completion of DEP_TASK_1
    router.notify_task_completion("DEP_TASK_1", agent_a.id)
    # Now DEP_TASK_2 should be routed (it was pending and its dependency is met)
    
    # Notify completion of DEP_TASK_2
    router.notify_task_completion("DEP_TASK_2", agent_a.id) # Assuming A handled it as it's analytics and can handle deps
    # Now DEP_TASK_3 should be routed (it was pending and its dependencies are met)

    print(f"\nFinal Agent A Load: {agent_a.current_load}")
    print(f"Final Agent B Load: {agent_b.current_load}")
    print(f"Final Agent C Load: {agent_c.current_load}")
    print(f"Final Agent D Load: {agent_d.current_load}")
    print(f"Pending tasks after all routing attempts: {len(router._pending_tasks)}")
    print(f"Completed tasks: {router._completed_tasks}")