"""
Orchestrator base classes for managing agents and patterns.

This module provides the foundation for orchestrating multiple agents
and patterns, including registration, coordination, and task delegation.
"""

from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional
from pydantic import BaseModel, Field

from pygeai_orchestration.core.base.agent import BaseAgent
from pygeai_orchestration.core.base.pattern import BasePattern, PatternResult


class OrchestratorConfig(BaseModel):
    """
    Configuration model for orchestrators.

    This Pydantic model defines orchestrator behavior including
    concurrency limits, timeouts, and retry policies.

    :param name: str - Unique identifier for the orchestrator.
    :param max_concurrent_tasks: int - Maximum number of tasks to run concurrently.
    :param timeout: Optional[float] - Global timeout in seconds. None means no timeout.
    :param retry_attempts: int - Number of retry attempts on task failure.
    :param metadata: Dict[str, Any] - Custom metadata for orchestrator configuration.
    """

    name: str = Field(..., description="Orchestrator name")
    max_concurrent_tasks: int = Field(5, ge=1, description="Max concurrent tasks")
    timeout: Optional[float] = Field(None, ge=0.0, description="Global timeout in seconds")
    retry_attempts: int = Field(3, ge=0, description="Number of retry attempts on failure")
    metadata: Dict[str, Any] = Field(default_factory=dict)


class BaseOrchestrator(ABC):
    """
    Abstract base class for orchestrator implementations.

    This class defines the interface for orchestrating multiple agents
    and patterns, enabling complex multi-agent workflows and pattern chaining.

    Orchestrators manage:
    - Agent registration and retrieval
    - Pattern registration and selection
    - Task delegation and coordination
    - Multi-agent collaboration

    Subclasses must implement:
    - orchestrate(): Execute a task using a specific pattern
    - coordinate_agents(): Coordinate multiple agents on a task
    """

    def __init__(self, config: OrchestratorConfig):
        """
        Initialize the orchestrator with configuration.

        :param config: OrchestratorConfig - Configuration for this orchestrator.
        """
        self.config = config
        self.name = config.name
        self._agents: Dict[str, BaseAgent] = {}
        self._patterns: Dict[str, BasePattern] = {}

    def register_agent(self, agent: BaseAgent) -> None:
        """
        Register an agent with this orchestrator.

        :param agent: BaseAgent - Agent to register.
        """
        self._agents[agent.name] = agent

    def register_pattern(self, pattern: BasePattern) -> None:
        """
        Register a pattern with this orchestrator.

        :param pattern: BasePattern - Pattern to register.
        """
        self._patterns[pattern.name] = pattern

    def get_agent(self, name: str) -> Optional[BaseAgent]:
        """
        Retrieve a registered agent by name.

        :param name: str - Agent name.
        :return: Optional[BaseAgent] - The agent if found, None otherwise.
        """
        return self._agents.get(name)

    def get_pattern(self, name: str) -> Optional[BasePattern]:
        """
        Retrieve a registered pattern by name.

        :param name: str - Pattern name.
        :return: Optional[BasePattern] - The pattern if found, None otherwise.
        """
        return self._patterns.get(name)

    def list_agents(self) -> List[str]:
        """
        List all registered agent names.

        :return: List[str] - Agent names.
        """
        return list(self._agents.keys())

    def list_patterns(self) -> List[str]:
        """
        List all registered pattern names.

        :return: List[str] - Pattern names.
        """
        return list(self._patterns.keys())

    @abstractmethod
    async def orchestrate(
        self, task: str, pattern: str, context: Optional[Dict[str, Any]] = None
    ) -> PatternResult:
        """
        Execute a task using a specified pattern.

        :param task: str - Task description or instruction.
        :param pattern: str - Name of the pattern to use.
        :param context: Optional[Dict[str, Any]] - Execution context.
        :return: PatternResult - Execution result.
        """
        pass

    @abstractmethod
    async def coordinate_agents(self, agents: List[str], task: str) -> Dict[str, Any]:
        """
        Coordinate multiple agents to complete a task.

        :param agents: List[str] - Names of agents to coordinate.
        :param task: str - Task description or instruction.
        :return: Dict[str, Any] - Coordination results.
        """
        pass

    def __repr__(self) -> str:
        return f"<{self.__class__.__name__}(name='{self.name}', agents={len(self._agents)}, patterns={len(self._patterns)})>"
