from typing import Literal

from pydantic import BaseModel

VisitEnvironment = Literal["outpatient", "emergency", "ward"]
AgeGroup = Literal["child", "youth", "middle_aged", "elderly"]
EducationLevel = Literal["primary_or_below", "secondary", "higher"]
Personality = Literal["calm", "anxious", "impatient", "cooperative", "suspicious"]


class ConfigOption(BaseModel):
    """训练配置选项：用于前端渲染单个可选项。"""

    value: str
    label: str
    description: str | None = None


class PatientConfig(BaseModel):
    """病人初始化配置：控制 AI 病人的就诊场景、年龄段、文化程度和性格。"""

    visit_environment: VisitEnvironment = "outpatient"
    age_group: AgeGroup = "youth"
    education_level: EducationLevel = "higher"
    personality: Personality = "calm"


class TrainingConfigOptionsResponse(BaseModel):
    """训练配置响应：返回推荐配置和全部可选项。"""

    case_id: int
    recommended: PatientConfig
    recommended_labels: dict[str, str]
    options: dict[str, list[ConfigOption]]


class TrainingConfigRecommendedResponse(BaseModel):
    """推荐训练配置响应：用于训练页进入时初始化默认病人信息。"""

    case_id: int
    recommended: PatientConfig
    recommended_labels: dict[str, str]
    options: dict[str, list[ConfigOption]]