from datetime import datetime, timezone
from typing import Any, Dict, List, Literal, Optional
from pydantic import BaseModel, Field


class UserMessage(BaseModel):
    """Message from user to agent through a communication channel."""
    session_id: str
    channel_id: str
    content: str
    attachments: List[str] = Field(default_factory=list, description="URLs or file paths")
    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
    metadata: Dict[str, Any] = Field(default_factory=dict)


class AgentMessage(BaseModel):
    """Message from agent to user(s) through one or more channels."""
    session_id: str
    target_channels: List[str] = Field(description="List of channel IDs to send to")
    content: str
    stream_chunk: bool = Field(default=False, description="True if this is a streaming chunk")
    done: bool = Field(default=False, description="True if streaming is complete")
    metadata: Dict[str, Any] = Field(default_factory=dict)


class ChannelStatus(BaseModel):
    """Status information about a communication channel."""
    channel_id: str
    channel_type: str = Field(description="Type: 'websocket', 'slack', 'telegram', etc.")
    connected: bool
    user_id: Optional[str] = None
    session_id: Optional[str] = None
    capabilities: Dict[str, bool] = Field(
        default_factory=lambda: {
            "streaming": False,
            "attachments": False,
            "markdown": False
        }
    )


# WebSocket-specific protocol extensions
class WebSocketAgentUserMessage(BaseModel):
    """WebSocket message: User → Backend (agent chat)"""
    type: Literal["agent_user_message"] = "agent_user_message"
    session_id: str
    content: str
    attachments: List[str] = Field(default_factory=list)


class WebSocketAgentChunk(BaseModel):
    """WebSocket message: Backend → User (streaming agent response)"""
    type: Literal["agent_chunk"] = "agent_chunk"
    session_id: str
    content: str
    done: bool = False
    metadata: Dict[str, Any] = Field(default_factory=dict)
