import random
from typing import List, Optional
from .agent import Agent


class AgentPool:
    """
    Feature 17: Agent Pooling — multiple instances of same agent.
    Feature 18: Load Balancing — distribute messages across pool.
    Feature 19: Failover — if one agent fails, use backup.
    """

    def __init__(self, agents: List[Agent], strategy: str = "round_robin"):
        """
        strategy: 'round_robin', 'random', 'least_used'
        """
        if not agents:
            raise ValueError("AgentPool requires at least one agent.")
        self.agents = agents
        self.strategy = strategy
        self._index = 0
        self._usage = {a.name: 0 for a in agents}

    def _pick_agent(self) -> Agent:
        """Feature 18: pick agent based on strategy."""
        if self.strategy == "round_robin":
            agent = self.agents[self._index % len(self.agents)]
            self._index += 1
            return agent
        elif self.strategy == "random":
            return random.choice(self.agents)
        elif self.strategy == "least_used":
            return min(self.agents, key=lambda a: self._usage[a.name])
        return self.agents[0]

    def complete(self, messages: list) -> str:
        """Feature 19: Failover — try next agent if one fails."""
        tried = set()
        for _ in range(len(self.agents)):
            agent = self._pick_agent()
            if agent.name in tried:
                continue
            tried.add(agent.name)
            try:
                response = agent.complete(messages)
                self._usage[agent.name] += 1
                return response
            except Exception as e:
                print(f"  [Pool] Agent '{agent.name}' failed: {e}. Trying next...")
                continue
        raise RuntimeError("All agents in pool failed.")

    def health(self) -> dict:
        return {a.name: a.health() for a in self.agents}

    def __repr__(self):
        return f"AgentPool(agents={[a.name for a in self.agents]}, strategy={self.strategy!r})"
