"""
Protocol definitions for swarm configuration and infrastructure.

These protocols define the interfaces for infrastructure components
like stream managers and metrics collectors, enabling dependency injection
and decoupling from specific implementations.

Design Principle: Behavior in Protocols, Not Enums
--------------------------------------------------
Instead of exposing state enums, protocols define behavior methods.
Swarm asks "is this done?" without knowing what "done" means.

Example:
    # Swarm code (state-agnostic):
    if instance.is_completed():
        ...

    # App implementation (owns the enum):
    def is_completed(self) -> bool:
        return self.state == WorkflowState.COMPLETED.value
"""

from typing import TYPE_CHECKING, Any, Dict, Optional, Protocol, Sequence

if TYPE_CHECKING:
    from swarm.message import SwarmMessage


class StreamManagerProtocol(Protocol):
    """
    Protocol for stream managers used throughout swarm.

    This interface allows swarm to work with any stream manager implementation
    that provides status publishing and error handling capabilities.
    """

    async def publish_status(self, status: str) -> None:
        """
        Publish a status message to the stream.

        Args:
            status: Status message to publish.
        """
        ...

    async def publish_error(self, message: str, code: int = 500) -> None:
        """
        Publish an error message to the stream.

        Args:
            message: Error message to publish.
            code: Error code (defaults to 500).
        """
        ...

    async def publish_message(self, content: Any, **kwargs: Any) -> None:
        """
        Publish a message to the stream.

        Args:
            content: Message content to publish.
            **kwargs: Additional fields for the message.
        """
        ...


class MetricsCollectorProtocol(Protocol):
    """
    Protocol for metrics collection in swarm operations.

    This interface allows swarm to work with any metrics collector
    that can record stage transitions.
    """

    def record_stage_transition(
        self, stage: str, status: str = "success", message_id: Optional[str] = None
    ) -> None:
        """
        Record a stage transition in the workflow.

        Args:
            stage: Name of the stage being entered.
            status: Status of the transition (default: "success").
            message_id: Optional message ID.
        """
        ...


class MessageStoreProtocol(Protocol):
    """
    Protocol for message persistence during streaming.

    This interface allows swarm to persist messages progressively
    as they complete, ensuring resilience against stream failures.
    The implementation handles conversion from SwarmMessage to
    application-specific message formats.
    """

    async def persist_messages(
        self,
        messages: Sequence["SwarmMessage"],
        conversation_id: str,
        context_variables: Dict[str, Any],
        designer_model: str,
        designer_fullname: str,
        user_id: Optional[str] = None,
        audience: Optional[str] = None,
        node_key: Optional[str] = None,
        is_client_message_override: Optional[bool] = None,
    ) -> None:
        """
        Persist messages to storage.

        Args:
            messages: Sequence of SwarmMessages to persist.
            conversation_id: ID of the conversation these messages belong to.
            context_variables: Context variables for the conversation.
            designer_model: Model of the agent.
            designer_fullname: Full name of the designer.
            user_id: Optional user ID for billing context.
            audience: Optional billing audience (e.g. "consumer", "studio").
            node_key: Optional workflow node key for billing attribution.
            is_client_message_override: When not None, forces is_client_message
                on assistant messages after transformer conversion.
        """
        ...


# =============================================================================
# WORKFLOW INSTANCE & NODE PROTOCOLS
# =============================================================================


class WorkflowInstanceProtocol(Protocol):
    """
    Protocol for workflow instance state queries.

    Exposes behavior methods instead of raw state values,
    allowing swarm to remain agnostic of specific state enums.
    """

    @property
    def id(self) -> Any:
        """Unique identifier for the workflow instance."""
        ...

    @property
    def workflow_id(self) -> Any:
        """ID of the workflow definition."""
        ...

    @property
    def state(self) -> str:
        """Current state as string (for logging/debugging)."""
        ...

    def is_completed(self) -> bool:
        """Check if workflow is in completed state."""
        ...

    def is_terminal(self) -> bool:
        """Check if workflow is in any terminal state (completed/failed/cancelled)."""
        ...


class NodeInstanceProtocol(Protocol):
    """
    Protocol for workflow node instance state queries.

    Exposes behavior methods instead of raw state values.
    """

    @property
    def id(self) -> Any:
        """Unique identifier for the node instance."""
        ...

    @property
    def node_id(self) -> str:
        """Node ID from the workflow definition."""
        ...

    @property
    def state(self) -> str:
        """Current state as string (for logging/debugging)."""
        ...

    @property
    def node_type(self) -> str:
        """Type of node (dialogue, agent, function, etc.)."""
        ...

    @property
    def configuration(self) -> Optional[Dict[str, Any]]:
        """Node configuration dictionary."""
        ...

    @property
    def metadata(self) -> Optional[Dict[str, Any]]:
        """Node metadata dictionary."""
        ...

    def is_terminal(self) -> bool:
        """Check if node is in terminal state (completed/failed/skipped/timed_out/cancelled)."""
        ...

    def is_running(self) -> bool:
        """Check if node is currently running."""
        ...


class WorkflowReadyNodeProtocol(Protocol):
    """Minimal shape needed to route and enqueue ready workflow nodes."""

    @property
    def id(self) -> Any:
        """Unique identifier for the node instance."""
        ...

    @property
    def node_id(self) -> str:
        """Node ID from the workflow definition."""
        ...

    @property
    def node_type(self) -> Any:
        """Type discriminator for routing decisions."""
        ...

    @property
    def configuration(self) -> Any:
        """Node configuration payload."""
        ...

    @property
    def metadata(self) -> Any:
        """Node metadata payload."""
        ...


class WorkflowRoutingDecisionProtocol(Protocol):
    """Minimal routing decision shape consumed by orchestrators."""

    @property
    def is_complete(self) -> bool:
        """Whether the workflow has no ready nodes."""
        ...

    @property
    def dialogue_node_instance_id(self) -> str | None:
        """Dialogue node to run next (if dialogue takes precedence)."""
        ...

    @property
    def autonomous_node_instance_ids(self) -> Sequence[str]:
        """Autonomous agent nodes ready for background enqueue."""
        ...

    @property
    def deterministic_node_instance_ids(self) -> Sequence[str]:
        """Deterministic nodes ready for background enqueue."""
        ...


# =============================================================================
# WORKFLOW SERVICE PROTOCOLS
# =============================================================================


class WorkflowNodeServiceProtocol(Protocol):
    """Protocol for workflow node operations."""

    async def get_ready_pending_nodes(
        self, instance_id: str
    ) -> Sequence[NodeInstanceProtocol]:
        """Get nodes ready for execution (dependencies satisfied)."""
        ...

    async def get_node_instances_by_workflow_instance_id(
        self, instance_id: str
    ) -> Sequence[NodeInstanceProtocol]:
        """Get all node instances for a workflow instance."""
        ...

    async def update_node_state(self, node_id: str, state: str) -> None:
        """Update a node's state."""
        ...

    async def store_task_id(self, node_id: str, task_id: str) -> None:
        """Store the background task ID for a node."""
        ...

    async def start_node(self, node_id: str) -> Any:
        """Mark a node as running (starting execution)."""
        ...


class WorkflowInstanceServiceProtocol(Protocol):
    """Protocol for workflow instance operations."""

    async def get_instance(
        self, instance_id: str
    ) -> Optional[WorkflowInstanceProtocol]:
        """Get a workflow instance by ID."""
        ...


class WorkflowServiceProtocol(Protocol):
    """Protocol for workflow orchestration operations."""

    # Properties use Any to allow duck typing with app's concrete service types
    instance_service: Any
    """Access to instance service (WorkflowInstanceServiceProtocol compatible)."""

    node_instance_service: Any
    """Access to node instance service (WorkflowNodeServiceProtocol compatible)."""

    async def get(self, workflow_id: str) -> Any:
        """Get a workflow definition by ID."""
        ...

    async def capture_workflow_snapshot(
        self, instance_id: str, trigger: str, description: Optional[str] = None
    ) -> Any:
        """Capture a snapshot of the workflow state."""
        ...

    async def get_ready_nodes_routing_decision(
        self,
        workflow_instance_id: str,
    ) -> tuple[Sequence[WorkflowReadyNodeProtocol], WorkflowRoutingDecisionProtocol]:
        """Return ready nodes and a routing decision for the current cycle."""
        ...

    async def complete_instance(
        self, instance_id: str, capture_snapshot: bool = True
    ) -> Any:
        """Mark workflow instance as completed."""
        ...


class PlanServiceProtocol(Protocol):
    """Protocol for plan orchestration."""

    async def handle_workflow_completion(
        self, workflow_instance_id: str, conversation_id: str
    ) -> None:
        """Handle completion of a workflow in a plan."""
        ...


# =============================================================================
# TASK BACKEND PROTOCOL
# =============================================================================


class TaskBackendProtocol(Protocol):
    """Protocol for background task execution."""

    async def enqueue(
        self,
        task_name: str,
        payload: Dict[str, Any],
    ) -> str:
        """
        Enqueue a task for background execution.

        Args:
            task_name: Name of the task to execute.
            payload: Task payload data.

        Returns:
            Task ID for tracking.
        """
        ...


# =============================================================================
# RETRY POLICY PROTOCOL
# =============================================================================


class RetryPolicyProtocol(Protocol):
    """Protocol for retry policy configuration."""

    @property
    def max_retries(self) -> int:
        """Maximum number of retry attempts."""
        ...

    @property
    def backoff(self) -> str:
        """Backoff strategy (constant, linear, exponential)."""
        ...

    @property
    def base_seconds(self) -> int:
        """Base delay in seconds."""
        ...

    @property
    def jitter(self) -> bool:
        """Whether to add jitter to delays."""
        ...

    @property
    def retry_on(self) -> Sequence[str]:
        """Error types that should trigger retries."""
        ...
