"""
Swarm protocols.

Defines the contracts for swarm coordination:
- SwarmParticipant: Any entity that can join a swarm
- SwarmCoordinator: Runs coordination logic on each participant

These protocols enable agents across tiers (cloud, edge, hardware)
to participate in self-coordinating swarms.

Example:
    class RobotAgent:
        @property
        def member_id(self) -> str:
            return "robot-buddy-01"

        @property
        def capabilities(self) -> list[str]:
            return ["drive", "look", "speak"]

        async def execute_task(self, task: dict, context: dict) -> TaskResult:
            # Execute task using robot-kit tools
            ...

        async def receive_message(self, message: dict) -> None:
            # Handle peer coordination messages
            ...
"""

from typing import Any, Dict, List, Optional, Protocol, runtime_checkable
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from dynabots_core.value_objects.task_result import TaskResult


@runtime_checkable
class SwarmParticipant(Protocol):
    """
    Protocol for entities that can participate in swarms.

    Any agent (cloud, edge, hardware) that implements this protocol
    can join and participate in swarm coordination.

    Required properties:
    - member_id: Unique identifier within the swarm
    - capabilities: What this participant can do

    Required methods:
    - execute_task: Perform a task assigned by swarm coordination
    - receive_message: Handle coordination messages from peers

    Optional methods:
    - get_status: Report current status for coordination
    - estimate_task: Estimate cost/time for a task (for auction allocation)
    """

    @property
    def member_id(self) -> str:
        """Unique identifier for this participant."""
        ...

    @property
    def capabilities(self) -> List[str]:
        """List of capability names this participant supports."""
        ...

    async def execute_task(
        self,
        task: Dict[str, Any],
        context: Dict[str, Any],
    ) -> "TaskResult":
        """
        Execute a task assigned by swarm coordination.

        Tasks are generated by the swarm to achieve objectives.
        Unlike agent process_task, swarm tasks include coordination context.

        Args:
            task: Task specification including:
                - task_id: Unique task identifier
                - action: What to do (capability or description)
                - objective_id: Which objective this serves
                - parameters: Task-specific parameters
            context: Execution context including:
                - swarm_id: Active swarm identifier
                - peer_results: Results from peer tasks
                - coordination_state: Current swarm state

        Returns:
            TaskResult with execution outcome.

        Example:
            async def execute_task(self, task, context):
                action = task["action"]
                if action == "patrol":
                    result = await self.tools.drive.execute({
                        "action": "forward",
                        "distance": task["parameters"]["distance"]
                    })
                    return TaskResult.success(task["task_id"], data=result)
        """
        ...

    async def receive_message(
        self,
        message: Dict[str, Any],
    ) -> None:
        """
        Handle coordination message from peers or coordinator.

        Messages enable self-organization without central control.
        Types include: task_bid, task_claim, heartbeat, status_update.

        Args:
            message: Message including:
                - type: Message type (bid, claim, heartbeat, etc.)
                - sender_id: Who sent this message
                - payload: Type-specific data

        Example:
            async def receive_message(self, message):
                if message["type"] == "task_bid":
                    # Consider bidding on this task
                    if self._can_do(message["payload"]["task"]):
                        await self.send_bid(message["payload"]["task_id"])
                elif message["type"] == "heartbeat":
                    # Record peer is alive
                    self._peer_heartbeats[message["sender_id"]] = time.time()
        """
        ...


@runtime_checkable
class SwarmParticipantWithStatus(Protocol):
    """
    Extended participant protocol with status reporting.

    Use this for participants that can report detailed status
    for intelligent task allocation and load balancing.
    """

    @property
    def member_id(self) -> str:
        ...

    @property
    def capabilities(self) -> List[str]:
        ...

    async def execute_task(
        self,
        task: Dict[str, Any],
        context: Dict[str, Any],
    ) -> "TaskResult":
        ...

    async def receive_message(
        self,
        message: Dict[str, Any],
    ) -> None:
        ...

    async def get_status(self) -> Dict[str, Any]:
        """
        Report current status for coordination decisions.

        Returns:
            Status dict including:
                - available: Whether can accept tasks
                - current_task: Task ID if busy
                - battery_level: For hardware (0.0-1.0)
                - location: For position-aware allocation
                - load: Current workload (0.0-1.0)

        Example:
            async def get_status(self):
                return {
                    "available": self.current_task is None,
                    "current_task": self.current_task,
                    "battery_level": await self.get_battery(),
                    "location": self.current_position,
                    "load": len(self.task_queue) / self.max_queue_size,
                }
        """
        ...

    async def estimate_task(
        self,
        task: Dict[str, Any],
    ) -> Dict[str, Any]:
        """
        Estimate cost/time for a task (for auction allocation).

        Used in distributed coordination to bid on tasks.

        Args:
            task: Task to estimate.

        Returns:
            Estimate dict including:
                - can_execute: Whether this participant can do it
                - estimated_duration_ms: Expected time
                - cost: Resource cost (battery, compute, etc.)
                - capability_score: How well-suited (0.0-1.0)

        Example:
            async def estimate_task(self, task):
                if not self._has_capabilities(task["required"]):
                    return {"can_execute": False}

                return {
                    "can_execute": True,
                    "estimated_duration_ms": self._estimate_duration(task),
                    "cost": self._calculate_cost(task),
                    "capability_score": self._score_match(task),
                }
        """
        ...


@runtime_checkable
class SwarmMessageBus(Protocol):
    """
    Protocol for swarm message transport.

    Abstracts the communication layer so swarms can run over:
    - Local function calls (testing, single-process)
    - WebSocket (edge networks)
    - MQTT (IoT/hardware)
    - gRPC (high-performance)
    """

    async def send(
        self,
        recipient_id: str,
        message: Dict[str, Any],
    ) -> bool:
        """
        Send message to specific participant.

        Returns True if delivered, False if recipient unavailable.
        """
        ...

    async def broadcast(
        self,
        message: Dict[str, Any],
        exclude: Optional[List[str]] = None,
    ) -> int:
        """
        Broadcast message to all participants.

        Args:
            message: Message to broadcast.
            exclude: List of member_ids to exclude.

        Returns:
            Number of participants that received the message.
        """
        ...

    async def subscribe(
        self,
        participant_id: str,
        handler: Any,  # Callable[[Dict[str, Any]], Awaitable[None]]
    ) -> None:
        """
        Subscribe participant to receive messages.

        Args:
            participant_id: Who is subscribing.
            handler: Async function to call with each message.
        """
        ...

    async def unsubscribe(
        self,
        participant_id: str,
    ) -> None:
        """Remove participant from message bus."""
        ...
