"""
WebSocket Protocol Schemas and Event Constants (Section 5, 6, 15).
Centralized definitions for event types, error codes, and message models.
"""

from enum import Enum
from typing import Optional, Any, Dict
from pydantic import BaseModel, Field


class WSEventType(str, Enum):
    CONNECTION_ACK = "connection.ack"
    PATIENT_MESSAGE = "patient.message"
    ASSISTANT_MESSAGE = "assistant.message"
    PROCESSING_STARTED = "processing.started"
    PROCESSING_COMPLETED = "processing.completed"
    TYPING_STARTED = "typing.started"
    TYPING_STOPPED = "typing.stopped"
    STATE_UPDATED = "state.updated"
    SAFETY_ALERT = "safety.alert"
    ONBOARDING_COMPLETED = "onboarding.completed"
    ERROR = "error"
    SESSION_ENDED = "session.ended"


class WSErrorCode(str, Enum):
    INVALID_MESSAGE = "INVALID_MESSAGE"
    MESSAGE_TOO_LARGE = "MESSAGE_TOO_LARGE"
    UNAUTHORIZED = "UNAUTHORIZED"
    SESSION_NOT_FOUND = "SESSION_NOT_FOUND"
    SESSION_EXPIRED = "SESSION_EXPIRED"
    PROCESSING_ERROR = "PROCESSING_ERROR"
    AI_TIMEOUT = "AI_TIMEOUT"
    AI_UNAVAILABLE = "AI_UNAVAILABLE"
    RATE_LIMITED = "RATE_LIMITED"
    INVALID_STATE = "INVALID_STATE"
    INTERNAL_ERROR = "INTERNAL_ERROR"


class WSConnectionAck(BaseModel):
    type: str = WSEventType.CONNECTION_ACK.value
    session_id: str
    status: str = "connected"


class WSPatientMessage(BaseModel):
    type: str = WSEventType.PATIENT_MESSAGE.value
    message_id: str = Field(..., description="Unique client message ID")
    text: str = Field(..., max_length=5000, description="Patient text message")


class WSAssistantMessage(BaseModel):
    type: str = WSEventType.ASSISTANT_MESSAGE.value
    message_id: str
    reply: str
    language: str = "en-IN"


class WSTypingEvent(BaseModel):
    type: str  # typing.started | typing.stopped


class WSProcessingStartedEvent(BaseModel):
    type: str = WSEventType.PROCESSING_STARTED.value


class WSProcessingCompletedEvent(BaseModel):
    type: str = WSEventType.PROCESSING_COMPLETED.value
    processing_time_ms: int


class WSSafetyAlertEvent(BaseModel):
    type: str = WSEventType.SAFETY_ALERT.value
    level: str = "emergency"  # emergency | crisis
    message: Optional[str] = None


class WSStateUpdatedEvent(BaseModel):
    type: str = WSEventType.STATE_UPDATED.value
    session_id: str


class WSOnboardingCompletedEvent(BaseModel):
    type: str = WSEventType.ONBOARDING_COMPLETED.value
    session_id: str


class WSErrorEvent(BaseModel):
    type: str = WSEventType.ERROR.value
    code: str
    message: str