from typing import Any, Dict, List, Literal, Optional, Union

from pydantic import BaseModel


class AuthMessage(BaseModel):
    """Authentication message (must be first message from client)"""
    type: Literal["auth"] = "auth"
    password: str
    confirm_password: Optional[str] = None  # Required only for initialization
    change_to_password: Optional[str] = None  # If provided, change password after auth
    confirm_new_password: Optional[str] = None  # Required if change_to_password is set

class AuthResponseMessage(BaseModel):
    """Authentication response from server"""
    type: Literal["auth_response"] = "auth_response"
    success: bool
    needs_confirmation: bool = False  # True if this is first-time setup
    password_changed: bool = False  # True if password was changed
    message: str

class SnapshotMessage(BaseModel):
    type: Literal["snapshot"] = "snapshot"
    store: str
    seq: int
    state: Dict[str, Any]

class PatchMessage(BaseModel):
    type: Literal["patch"] = "patch"
    store: str
    seq: int
    patch: List[Dict[str, Any]]

class HelloMessage(BaseModel):
    type: Literal["hello"] = "hello"
    seqs: Dict[str, int]

# Union type for all messages from backend to frontend
BackendMessage = Union[SnapshotMessage, PatchMessage, AuthResponseMessage]

# Union type for all messages from frontend to backend
FrontendMessage = Union[AuthMessage, HelloMessage, PatchMessage]