"""
Base WebSocket message schemas.

Provides the foundation for all WebSocket message types including:
- Base message class with type and timestamp
- Error message structure
- Common enums and types

Usage:
    from lightwave.schema.pydantic.contracts.websockets.base import WSMessage, WSErrorMessage

    class MyCustomMessage(WSMessage):
        type: Literal["custom.message"] = "custom.message"
        data: str
"""

from datetime import datetime
from typing import Any, Literal

from pydantic import Field

from lightwave.schema.pydantic.core.base import LightwaveBaseSchema

# =============================================================================
# Enums as Literals
# =============================================================================

MessageDirection = Literal["inbound", "outbound", "both"]

WSChannel = Literal[
    "ai_chat",
    "notifications",
    "sync",
    "presence",
    "collaboration",
    "system",
]

WSMessageType = Literal[
    # Chat
    "chat.message",
    "chat.response",
    "chat.stream",
    "chat.typing",
    "chat.stop",
    "chat.feedback",
    "chat.error",
    "chat.complete",
    # Notifications
    "notification.new",
    "notification.update",
    "notification.batch",
    "notification.ack",
    "notification.dismiss",
    "notification.mark_read",
    # Sync
    "sync.subscribe",
    "sync.unsubscribe",
    "sync.ack",
    "sync.checkpoint",
    "sync.checkpoint_ack",
    "sync.update",
    "sync.delete",
    "sync.batch",
    "sync.conflict",
    # Presence
    "presence.update",
    "presence.typing",
    "presence.away",
    "presence.subscribe",
    "presence.user_joined",
    "presence.user_left",
    "presence.user_typing",
    "presence.user_away",
    "presence.roster",
    # Collaboration
    "collab.operation",
    "collab.cursor",
    "collab.selection",
    "collab.lock_request",
    "collab.lock_granted",
    "collab.lock_denied",
    "collab.unlock",
    "collab.conflict",
    # System
    "system.ping",
    "system.pong",
    "system.health",
    "system.maintenance",
    "system.broadcast",
    # Error
    "error",
]

WSErrorCode = Literal[
    # Connection errors
    "WS_AUTH_FAILED",
    "WS_AUTH_EXPIRED",
    "WS_FORBIDDEN",
    "WS_RATE_LIMITED",
    "WS_CONNECTION_LIMIT",
    # Message errors
    "WS_INVALID_MESSAGE",
    "WS_UNKNOWN_TYPE",
    "WS_PAYLOAD_TOO_LARGE",
    "WS_MISSING_FIELD",
    # Channel errors
    "WS_CHANNEL_NOT_FOUND",
    "WS_CHANNEL_CLOSED",
    # Processing errors
    "WS_INTERNAL_ERROR",
    "WS_TIMEOUT",
    "WS_CONFLICT",
]


# =============================================================================
# Base Message Classes
# =============================================================================


class WSMessage(LightwaveBaseSchema):
    """
    Base class for all WebSocket messages.

    All WebSocket messages have a type and timestamp. Subclasses
    should override the `type` field with a Literal type.

    Example:
        class ChatMessage(WSMessage):
            type: Literal["chat.message"] = "chat.message"
            content: str
            conversation_id: str
    """

    type: str = Field(..., description="Message type identifier")
    timestamp: datetime = Field(
        default_factory=datetime.utcnow,
        description="Message timestamp (UTC)",
    )
    request_id: str | None = Field(
        None,
        description="Optional request ID for request-response correlation",
    )

    def to_ws_json(self) -> str:
        """Serialize to JSON string for WebSocket transmission."""
        return self.model_dump_json(exclude_none=True)


class WSErrorMessage(WSMessage):
    """
    Standard WebSocket error message.

    Sent when an error occurs processing a message or in the connection.

    Example:
        error = WSErrorMessage(
            error_code="WS_AUTH_EXPIRED",
            error_message="Your session has expired",
            recoverable=True,
        )
        await websocket.send(error.to_ws_json())
    """

    type: Literal["error"] = "error"
    error_code: WSErrorCode = Field(..., description="Machine-readable error code")
    error_message: str = Field(..., description="Human-readable error message")
    recoverable: bool = Field(
        True,
        description="Whether client can recover (retry/reconnect)",
    )
    details: dict[str, Any] | None = Field(
        None,
        description="Additional error context",
    )
    original_type: str | None = Field(
        None,
        description="Type of message that caused the error",
    )


class WSAckMessage(WSMessage):
    """
    Generic acknowledgment message.

    Used to confirm receipt of messages that require acknowledgment.
    """

    type: Literal["ack"] = "ack"
    ack_type: str = Field(..., description="Type of message being acknowledged")
    ack_id: str = Field(..., description="ID of message being acknowledged")
    success: bool = Field(True, description="Whether processing succeeded")


# =============================================================================
# Message Envelope (for routing)
# =============================================================================


class WSEnvelope(LightwaveBaseSchema):
    """
    Message envelope for routing and metadata.

    Used internally for message routing and doesn't go over the wire.
    """

    channel: WSChannel = Field(..., description="Target channel")
    user_id: str = Field(..., description="Sender/recipient user ID")
    tenant_id: str = Field(..., description="Tenant context")
    team_id: str | None = Field(None, description="Team context if applicable")
    message: WSMessage = Field(..., description="The actual message")
    received_at: datetime = Field(
        default_factory=datetime.utcnow,
        description="When message was received by server",
    )
