"""
Message Data Model - Structured messages for agent-to-agent communication.

This module defines the EXACT schema for how agents communicate within a workspace.
No loose natural language - every message has a type, structure, and routing rules.
"""

from dataclasses import dataclass, field
from enum import Enum
from typing import Any, Optional
from datetime import datetime
import uuid


class MessageType(Enum):
    """
    Structured message types for agent communication.
    
    Each type has specific semantics and routing rules.
    """
    # Core interaction types
    SUGGESTION = "suggestion"  # Agent proposes an idea or approach
    OBSERVATION = "observation"  # Agent shares a fact or finding
    PLAN_FRAGMENT = "plan_fragment"  # Agent contributes to the plan
    REQUEST = "request"  # Agent requests information or action
    RESPONSE = "response"  # Agent responds to a request
    CRITIQUE = "critique"  # Agent provides critical feedback
    VOTE = "vote"  # Agent votes on a proposal
    TOOL_CALL = "tool_call"  # Agent invokes a tool or capability
    FINAL_OUTPUT_PROPOSAL = "final_output_proposal"  # Agent proposes final answer
    
    # Meta-communication (workspace management)
    SYSTEM = "system"  # Workspace/orchestrator message
    STATUS_UPDATE = "status_update"  # Agent reports progress
    CLARIFICATION = "clarification"  # Agent seeks clarification
    ACKNOWLEDGMENT = "acknowledgment"  # Agent acknowledges a message


class MessagePriority(Enum):
    """Message priority for scheduling."""
    CRITICAL = 1  # Must be processed immediately
    HIGH = 2  # Should be processed soon
    NORMAL = 3  # Standard priority
    LOW = 4  # Can be deferred


@dataclass
class MessageMetadata:
    """Metadata about a message."""
    message_id: str = field(default_factory=lambda: str(uuid.uuid4()))
    timestamp: datetime = field(default_factory=datetime.now)
    priority: MessagePriority = MessagePriority.NORMAL
    
    # Routing
    reply_to: Optional[str] = None  # Message ID this is replying to
    thread_id: Optional[str] = None  # Conversation thread
    requires_response: bool = False  # Does this message require a response?
    
    # Context
    round_number: int = 0  # Which round of communication
    phase: Optional[str] = None  # Which phase of the protocol
    
    # Constraints
    token_count: int = 0  # Approximate token count
    max_token_limit: Optional[int] = None  # Max tokens allowed for responses
    
    # Custom fields
    custom: dict[str, Any] = field(default_factory=dict)


@dataclass
class StructuredMessage:
    """
    A structured message between agents.
    
    This is the core unit of communication in the workspace.
    Every agent interaction is expressed as a StructuredMessage.
    """
    # Core fields (required)
    message_type: MessageType
    sender_id: str
    sender_name: str
    content: str  # The actual message content
    
    # Routing (optional)
    recipient_ids: list[str] = field(default_factory=list)  # Empty = broadcast
    recipient_names: list[str] = field(default_factory=list)
    
    # Metadata
    metadata: MessageMetadata = field(default_factory=MessageMetadata)
    
    # Structured data (type-specific)
    structured_data: dict[str, Any] = field(default_factory=dict)
    
    # State
    processed: bool = False
    responses: list[str] = field(default_factory=list)  # Response message IDs
    
    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary for serialization."""
        return {
            "message_id": self.metadata.message_id,
            "message_type": self.message_type.value,
            "sender_id": self.sender_id,
            "sender_name": self.sender_name,
            "content": self.content,
            "recipient_ids": self.recipient_ids,
            "recipient_names": self.recipient_names,
            "timestamp": self.metadata.timestamp.isoformat(),
            "priority": self.metadata.priority.value,
            "reply_to": self.metadata.reply_to,
            "thread_id": self.metadata.thread_id,
            "requires_response": self.metadata.requires_response,
            "round_number": self.metadata.round_number,
            "phase": self.metadata.phase,
            "token_count": self.metadata.token_count,
            "processed": self.processed,
            "structured_data": self.structured_data,
        }
    
    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "StructuredMessage":
        """Create from dictionary."""
        metadata = MessageMetadata(
            message_id=data.get("message_id", str(uuid.uuid4())),
            timestamp=datetime.fromisoformat(data.get("timestamp", datetime.now().isoformat())),
            priority=MessagePriority(data.get("priority", MessagePriority.NORMAL.value)),
            reply_to=data.get("reply_to"),
            thread_id=data.get("thread_id"),
            requires_response=data.get("requires_response", False),
            round_number=data.get("round_number", 0),
            phase=data.get("phase"),
            token_count=data.get("token_count", 0),
        )
        
        return cls(
            message_type=MessageType(data["message_type"]),
            sender_id=data["sender_id"],
            sender_name=data["sender_name"],
            content=data["content"],
            recipient_ids=data.get("recipient_ids", []),
            recipient_names=data.get("recipient_names", []),
            metadata=metadata,
            structured_data=data.get("structured_data", {}),
            processed=data.get("processed", False),
        )
    
    def is_broadcast(self) -> bool:
        """Check if this is a broadcast message."""
        return len(self.recipient_ids) == 0
    
    def is_directed_to(self, agent_id: str) -> bool:
        """Check if this message is directed to a specific agent."""
        return self.is_broadcast() or agent_id in self.recipient_ids
    
    def requires_response_from(self, agent_id: str) -> bool:
        """Check if this message requires a response from a specific agent."""
        return self.metadata.requires_response and self.is_directed_to(agent_id)


# Type-specific structured data schemas
# These define what goes in the structured_data field for each message type

@dataclass
class SuggestionData:
    """Structured data for SUGGESTION messages."""
    suggestion_id: str
    rationale: str
    confidence: float = 0.0  # 0.0 to 1.0
    alternatives: list[str] = field(default_factory=list)


@dataclass
class VoteData:
    """Structured data for VOTE messages."""
    vote_id: str
    option: str
    rationale: Optional[str] = None
    confidence: float = 1.0


@dataclass
class CritiqueData:
    """Structured data for CRITIQUE messages."""
    critique_id: str
    target_message_id: str  # What is being critiqued
    strengths: list[str] = field(default_factory=list)
    weaknesses: list[str] = field(default_factory=list)
    suggestions: list[str] = field(default_factory=list)
    severity: str = "info"  # info, warning, critical


@dataclass
class PlanFragmentData:
    """Structured data for PLAN_FRAGMENT messages."""
    fragment_id: str
    step_number: Optional[int] = None
    dependencies: list[str] = field(default_factory=list)  # Other fragment IDs
    estimated_effort: Optional[str] = None
    assigned_to: Optional[str] = None


@dataclass
class FinalOutputProposalData:
    """Structured data for FINAL_OUTPUT_PROPOSAL messages."""
    proposal_id: str
    confidence: float = 0.0
    supporting_evidence: list[str] = field(default_factory=list)
    caveats: list[str] = field(default_factory=list)


# Helper functions for creating structured messages

def create_suggestion(
    sender_id: str,
    sender_name: str,
    content: str,
    rationale: str,
    confidence: float = 0.5,
    recipient_ids: Optional[list[str]] = None,
    round_number: int = 0,
) -> StructuredMessage:
    """Create a SUGGESTION message."""
    suggestion_id = str(uuid.uuid4())
    
    metadata = MessageMetadata(
        round_number=round_number,
        priority=MessagePriority.NORMAL,
    )
    
    structured_data = {
        "suggestion_id": suggestion_id,
        "rationale": rationale,
        "confidence": confidence,
    }
    
    return StructuredMessage(
        message_type=MessageType.SUGGESTION,
        sender_id=sender_id,
        sender_name=sender_name,
        content=content,
        recipient_ids=recipient_ids or [],
        metadata=metadata,
        structured_data=structured_data,
    )


def create_critique(
    sender_id: str,
    sender_name: str,
    target_message_id: str,
    content: str,
    strengths: list[str],
    weaknesses: list[str],
    suggestions: list[str],
    round_number: int = 0,
) -> StructuredMessage:
    """Create a CRITIQUE message."""
    critique_id = str(uuid.uuid4())
    
    metadata = MessageMetadata(
        round_number=round_number,
        reply_to=target_message_id,
        priority=MessagePriority.HIGH,
    )
    
    structured_data = {
        "critique_id": critique_id,
        "target_message_id": target_message_id,
        "strengths": strengths,
        "weaknesses": weaknesses,
        "suggestions": suggestions,
    }
    
    return StructuredMessage(
        message_type=MessageType.CRITIQUE,
        sender_id=sender_id,
        sender_name=sender_name,
        content=content,
        metadata=metadata,
        structured_data=structured_data,
    )


def create_vote(
    sender_id: str,
    sender_name: str,
    option: str,
    rationale: Optional[str] = None,
    confidence: float = 1.0,
    round_number: int = 0,
) -> StructuredMessage:
    """Create a VOTE message."""
    vote_id = str(uuid.uuid4())
    
    metadata = MessageMetadata(
        round_number=round_number,
        priority=MessagePriority.HIGH,
    )
    
    structured_data = {
        "vote_id": vote_id,
        "option": option,
        "rationale": rationale,
        "confidence": confidence,
    }
    
    return StructuredMessage(
        message_type=MessageType.VOTE,
        sender_id=sender_id,
        sender_name=sender_name,
        content=f"Vote for: {option}" + (f" - {rationale}" if rationale else ""),
        metadata=metadata,
        structured_data=structured_data,
    )


def create_final_output_proposal(
    sender_id: str,
    sender_name: str,
    content: str,
    confidence: float,
    supporting_evidence: list[str],
    round_number: int = 0,
) -> StructuredMessage:
    """Create a FINAL_OUTPUT_PROPOSAL message."""
    proposal_id = str(uuid.uuid4())
    
    metadata = MessageMetadata(
        round_number=round_number,
        priority=MessagePriority.CRITICAL,
    )
    
    structured_data = {
        "proposal_id": proposal_id,
        "confidence": confidence,
        "supporting_evidence": supporting_evidence,
    }
    
    return StructuredMessage(
        message_type=MessageType.FINAL_OUTPUT_PROPOSAL,
        sender_id=sender_id,
        sender_name=sender_name,
        content=content,
        metadata=metadata,
        structured_data=structured_data,
    )

