"""
AI Chat WebSocket message schemas.

Provides typed messages for real-time AI chat conversations including:
- User messages to AI
- AI responses (complete and streaming)
- Typing indicators
- Stop/cancel requests
- Error handling

Usage:
    from lightwave.schema.pydantic.contracts.websockets.chat import (
        ChatMessageInbound,
        ChatResponseOutbound,
        ChatStreamOutbound,
    )

    # Parse user message
    msg = ChatMessageInbound.model_validate(raw_data)

    # Send streaming response
    for chunk in ai_stream:
        stream_msg = ChatStreamOutbound(
            delta=chunk,
            conversation_id=msg.conversation_id,
            message_id=response_id,
            is_final=False,
        )
        await websocket.send(stream_msg.to_ws_json())
"""

from typing import Any, Literal

from pydantic import Field

from lightwave.schema.pydantic.contracts.websockets.base import WSErrorCode, WSMessage

# =============================================================================
# Attachment Schema
# =============================================================================


class ChatAttachment(WSMessage):
    """Attachment in a chat message."""

    type: Literal["image", "file", "code"] = Field(..., description="Attachment type")
    url: str = Field(..., description="URL to the attachment")
    name: str = Field(..., description="Display name")
    mime_type: str | None = Field(None, description="MIME type if known")
    size_bytes: int | None = Field(None, description="File size in bytes")


# =============================================================================
# Inbound Messages (Client -> Server)
# =============================================================================


class ChatMessageInbound(WSMessage):
    """
    User sends a message to AI.

    Example:
        {
            "type": "chat.message",
            "content": "Hello, can you help me?",
            "conversation_id": "abc-123",
            "context": {"page": "/dashboard"}
        }
    """

    type: Literal["chat.message"] = "chat.message"
    content: str = Field(..., max_length=10000, description="Message content")
    conversation_id: str = Field(..., description="Conversation UUID")
    context: dict[str, Any] | None = Field(
        None,
        description="Additional context for AI (current page, selection, etc.)",
    )
    attachments: list[ChatAttachment] | None = Field(
        None,
        description="File/image attachments",
    )
    model_preference: str | None = Field(
        None,
        description="Preferred AI model (if allowed)",
    )


class ChatTypingMessage(WSMessage):
    """
    Typing indicator (bidirectional).

    Sent by client when user is typing, and by server when AI is "thinking".
    """

    type: Literal["chat.typing"] = "chat.typing"
    conversation_id: str = Field(..., description="Conversation UUID")
    is_typing: bool = Field(..., description="Whether typing is in progress")


class ChatStopInbound(WSMessage):
    """
    User requests to stop AI generation.

    Sent when user clicks "Stop" during streaming response.
    """

    type: Literal["chat.stop"] = "chat.stop"
    conversation_id: str = Field(..., description="Conversation UUID")
    message_id: str | None = Field(
        None,
        description="Specific message to stop (if multiple in progress)",
    )


class ChatFeedbackInbound(WSMessage):
    """
    User provides feedback on AI response.

    Example:
        {
            "type": "chat.feedback",
            "conversation_id": "abc-123",
            "message_id": "msg-456",
            "rating": "positive",
            "comment": "Very helpful!"
        }
    """

    type: Literal["chat.feedback"] = "chat.feedback"
    conversation_id: str = Field(..., description="Conversation UUID")
    message_id: str = Field(..., description="Message being rated")
    rating: Literal["positive", "negative", "neutral"] = Field(..., description="Feedback rating")
    comment: str | None = Field(None, max_length=1000, description="Optional comment")
    categories: list[str] | None = Field(
        None,
        description="Feedback categories (e.g., 'helpful', 'accurate', 'creative')",
    )


# =============================================================================
# Outbound Messages (Server -> Client)
# =============================================================================


class ChatUsageStats(WSMessage):
    """Token usage statistics for a response."""

    input_tokens: int = Field(0, ge=0, description="Input tokens used")
    output_tokens: int = Field(0, ge=0, description="Output tokens generated")
    total_tokens: int = Field(0, ge=0, description="Total tokens")
    cached_tokens: int = Field(0, ge=0, description="Tokens from cache")


class ChatResponseOutbound(WSMessage):
    """
    Complete AI response (non-streaming).

    Sent when streaming is disabled or as final message after stream.
    """

    type: Literal["chat.response"] = "chat.response"
    content: str = Field(..., description="Full response content")
    conversation_id: str = Field(..., description="Conversation UUID")
    message_id: str = Field(..., description="Unique message ID")
    model: str | None = Field(None, description="Model that generated response")
    usage: ChatUsageStats | None = Field(None, description="Token usage")
    metadata: dict[str, Any] | None = Field(
        None,
        description="Additional response metadata",
    )


class ChatStreamOutbound(WSMessage):
    """
    Streaming AI response chunk.

    Sent multiple times during response generation.
    Client should accumulate `delta` values to build full response.
    """

    type: Literal["chat.stream"] = "chat.stream"
    delta: str = Field(..., description="New content chunk")
    conversation_id: str = Field(..., description="Conversation UUID")
    message_id: str = Field(..., description="Message being streamed")
    is_final: bool = Field(..., description="Whether this is the last chunk")
    index: int = Field(0, ge=0, description="Chunk index for ordering")


class ChatCompleteOutbound(WSMessage):
    """
    Signal that chat response is complete.

    Sent after all stream chunks or after complete response.
    Includes final statistics and metadata.
    """

    type: Literal["chat.complete"] = "chat.complete"
    conversation_id: str = Field(..., description="Conversation UUID")
    message_id: str = Field(..., description="Completed message ID")
    usage: ChatUsageStats | None = Field(None, description="Final token usage")
    duration_ms: int | None = Field(None, description="Response generation time")
    stopped: bool = Field(False, description="Whether stopped by user request")


class ChatErrorOutbound(WSMessage):
    """
    Error in chat processing.

    Sent when AI fails to generate response or other error occurs.
    """

    type: Literal["chat.error"] = "chat.error"
    error_code: WSErrorCode | str = Field(..., description="Error code")
    error_message: str = Field(..., description="Human-readable error")
    conversation_id: str | None = Field(None, description="Conversation if known")
    message_id: str | None = Field(None, description="Message that failed")
    recoverable: bool = Field(True, description="Whether client can retry")
    retry_after_ms: int | None = Field(
        None,
        description="Suggested retry delay (for rate limits)",
    )
