"""
Protocol-Based Agent Communication Module.

This module defines explicit message protocols for agent-to-agent communication,
replacing implicit state passing with formal message intents.

Following LangGraph best practices for multi-agent systems:
- Agents communicate via structured messages (not just state dicts)
- Message intents are explicit (TASK_REQUEST, TASK_RESPONSE, etc.)
- Human-in-the-loop checkpoints are supported
"""

from datetime import UTC, datetime
from enum import Enum
from typing import Any, TypedDict

from pydantic import BaseModel, Field


class MessageIntent(str, Enum):
    """
    Explicit message intents for agent communication.

    Following protocol-based communication patterns from:
    https://bix-tech.com/agent-to-agent-communication-with-langgraph-protocol-based-workflows-a-practical-guide/
    """

    # Task Execution
    TASK_REQUEST = "task_request"  # "Please complete this subtask"
    TASK_RESPONSE = "task_response"  # "Here's my output"

    # Control Flow
    ESCALATION = "escalation"  # "I need human help"
    CLARIFICATION_REQUEST = "clarification_request"  # "I need more info"
    APPROVAL_REQUEST = "approval_request"  # "Please approve this output"

    # Error Handling
    ERROR = "error"  # "I failed because..."
    RETRY_REQUEST = "retry_request"  # "Please retry this task"

    # Completion
    TASK_COMPLETE = "task_complete"  # "Task finished successfully"
    PIPELINE_COMPLETE = "pipeline_complete"  # "All tasks finished"


class MessagePriority(str, Enum):
    """Message priority levels for routing decisions."""

    LOW = "low"
    NORMAL = "normal"
    HIGH = "high"
    CRITICAL = "critical"


class AgentMessage(BaseModel):
    """
    Structured message between agents.

    This replaces implicit state passing with explicit message protocols.
    """

    # Message identification
    id: str = Field(default_factory=lambda: datetime.now(UTC).isoformat())

    # Protocol fields
    intent: MessageIntent
    priority: MessagePriority = MessagePriority.NORMAL

    # Communication metadata
    sender: str  # TeamRole value
    recipient: str | None = None  # TeamRole value or None for broadcast

    # Content
    content: str
    metadata: dict[str, Any] = Field(default_factory=dict)

    # Timestamps
    created_at: str = Field(default_factory=lambda: datetime.now(UTC).isoformat())
    processed_at: str | None = None

    # Status
    status: str = "pending"  # pending, processing, processed, failed

    def mark_processed(self) -> "AgentMessage":
        """Mark message as processed."""
        self.processed_at = datetime.now(UTC).isoformat()
        self.status = "processed"
        return self

    def mark_failed(self, error: str) -> "AgentMessage":
        """Mark message as failed."""
        self.metadata["error"] = error
        self.status = "failed"
        return self


class MessageQueue(TypedDict):
    """
    Queue of messages in agent state.

    This replaces the implicit history passing with explicit message queues.
    """

    messages: list[AgentMessage]
    pending_count: int
    processed_count: int


def merge_messages(a: list[AgentMessage], b: list[AgentMessage]) -> list[AgentMessage]:
    """
    Reducer function for merging message lists.

    Combines messages from multiple agents while maintaining order.
    """
    return a + b


def create_task_request_message(
    sender: str,
    recipient: str,
    content: str,
    metadata: dict | None = None,
) -> AgentMessage:
    """
    Factory function to create a task request message.

    Use this when an agent needs to delegate work to another agent.
    """
    return AgentMessage(
        intent=MessageIntent.TASK_REQUEST,
        sender=sender,
        recipient=recipient,
        content=content,
        metadata=metadata or {},
        priority=MessagePriority.HIGH,
    )


def create_task_response_message(
    sender: str,
    content: str,
    status: str = "completed",
    metadata: dict | None = None,
) -> AgentMessage:
    """
    Factory function to create a task response message.

    Use this when an agent completes a task and returns results.
    """
    return AgentMessage(
        intent=MessageIntent.TASK_RESPONSE,
        sender=sender,
        content=content,
        metadata=metadata or {"status": status},
        priority=MessagePriority.NORMAL,
    )


def create_escalation_message(
    sender: str,
    content: str,
    reason: str,
    metadata: dict | None = None,
) -> AgentMessage:
    """
    Factory function to create an escalation message.

    Use this when an agent needs human intervention or cannot proceed.
    """
    return AgentMessage(
        intent=MessageIntent.ESCALATION,
        sender=sender,
        content=content,
        metadata=metadata or {"reason": reason},
        priority=MessagePriority.CRITICAL,
    )


def create_error_message(
    sender: str,
    error: str,
    context: str | None = None,
    metadata: dict | None = None,
) -> AgentMessage:
    """
    Factory function to create an error message.

    Use this when an agent encounters an error it cannot recover from.
    """
    return AgentMessage(
        intent=MessageIntent.ERROR,
        sender=sender,
        content=error,
        metadata=metadata or {"context": context},
        priority=MessagePriority.HIGH,
    )


def create_clarification_message(
    sender: str,
    question: str,
    metadata: dict | None = None,
) -> AgentMessage:
    """
    Factory function to create a clarification request message.

    Use this when an agent needs more information to proceed.
    """
    return AgentMessage(
        intent=MessageIntent.CLARIFICATION_REQUEST,
        sender=sender,
        content=question,
        metadata=metadata or {},
        priority=MessagePriority.NORMAL,
    )


def get_pending_messages(messages: list[AgentMessage]) -> list[AgentMessage]:
    """Filter messages that are still pending."""
    return [m for m in messages if m.status == "pending"]


def get_messages_by_intent(
    messages: list[AgentMessage], intent: MessageIntent
) -> list[AgentMessage]:
    """Filter messages by intent type."""
    return [m for m in messages if m.intent == intent]


def get_messages_for_recipient(
    messages: list[AgentMessage], recipient: str
) -> list[AgentMessage]:
    """Filter messages for a specific recipient."""
    return [m for m in messages if m.recipient == recipient or m.recipient is None]