"""
Schemas Pydantic du protocole WebSocket /ws/edit.

Deux messages client -> serveur (protocole a 2 phases, cf. decision UX
"apercu du masque avant generation") :
  1. LocateRequest  -> le serveur repond MaskReadyEvent (ou ErrorEvent)
  2. ConfirmRequest -> le serveur repond une serie de ProgressEvent puis un
                        ResultEvent (ou ErrorEvent)

Une connexion WebSocket = une session d'edition (l'etat image+masque entre
les deux phases vit le temps de la connexion, cote serveur).
"""

from typing import List, Literal, Optional

from pydantic import BaseModel


class LocateRequest(BaseModel):
    type: Literal["locate"] = "locate"
    image_b64: str
    text_query: str


class ConfirmRequest(BaseModel):
    type: Literal["confirm"] = "confirm"
    prompt: str
    negative_prompt: str = ""
    steps: int = 30
    guidance_scale: float = 7.5
    seed: int = 42


class MaskReadyEvent(BaseModel):
    type: Literal["mask_ready"] = "mask_ready"
    mask_b64: str
    box: List[float]
    clip_confidence: float


class ProgressEvent(BaseModel):
    type: Literal["progress"] = "progress"
    stage: Literal["diffusion"] = "diffusion"
    step: int
    total: int


class ResultEvent(BaseModel):
    type: Literal["result"] = "result"
    image_b64: str


class ErrorEvent(BaseModel):
    type: Literal["error"] = "error"
    message: str