"""WebSocket message models for AiCippy real-time communication.

Provides typed dataclasses for WebSocket message serialization/deserialization,
stream tokens, and connection state management.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from enum import StrEnum
from typing import Any


class ConnectionState(StrEnum):
    """WebSocket connection lifecycle states."""

    DISCONNECTED = "disconnected"
    CONNECTING = "connecting"
    CONNECTED = "connected"
    RECONNECTING = "reconnecting"


@dataclass(frozen=True, slots=True)
class StreamToken:
    """A single token received from the WebSocket stream.

    Attributes:
        text: The token text content.
        done: Whether this is the final token in the stream.
        input_tokens: Total input tokens consumed (populated on done=True).
        output_tokens: Total output tokens generated (populated on done=True).

    """

    text: str
    done: bool
    input_tokens: int = 0
    output_tokens: int = 0


@dataclass(slots=True)
class ChatRequest:
    """Outbound chat request to be sent over WebSocket.

    Attributes:
        message: The user message text.
        model: Model identifier (e.g. 'llama', 'opus').
        system_prompt: System prompt for the model.
        conversation_history: Prior conversation turns.
        correlation_id: Unique ID to correlate request with streamed response.

    """

    message: str
    model: str
    system_prompt: str
    conversation_history: list[dict[str, Any]] = field(default_factory=list)
    correlation_id: str = ""

    def to_wire(self) -> dict[str, Any]:
        """Serialize to the wire format expected by API Gateway."""
        return {
            "type": "chat",
            "payload": {
                "message": self.message,
                "model": self.model,
                "system_prompt": self.system_prompt,
                "conversation_history": self.conversation_history,
            },
            "correlation_id": self.correlation_id,
        }


@dataclass(slots=True)
class WebSocketMessage:
    """Generic WebSocket message envelope (inbound or outbound).

    Attributes:
        type: Message type discriminator (e.g. 'chat', 'token', 'ping', 'pong', 'error').
        payload: Type-specific payload data.
        correlation_id: Correlation ID linking request to response stream.

    """

    type: str
    payload: dict[str, Any] = field(default_factory=dict)
    correlation_id: str = ""

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> WebSocketMessage:
        """Deserialize from a parsed JSON dictionary."""
        return cls(
            type=data.get("type", "unknown"),
            payload=data.get("payload", {}),
            correlation_id=data.get("correlation_id", ""),
        )

    def to_dict(self) -> dict[str, Any]:
        """Serialize to a JSON-compatible dictionary."""
        return {
            "type": self.type,
            "payload": self.payload,
            "correlation_id": self.correlation_id,
        }
