from typing import Any, Literal, Optional
from pydantic import BaseModel, Field


class PatientContext(BaseModel):
    id: str = Field(..., min_length=1)
    name: Optional[str] = None
    age: Optional[int] = None
    gender: Optional[str] = None
    abha_id: Optional[str] = None
    phone: Optional[str] = None
    insurance_scheme: Optional[str] = None
    dietary_habits: Optional[str] = None


class CategoryContext(BaseModel):
    type: Literal["general", "speciality"]
    code: str = Field(..., min_length=1)
    name: str = Field(..., min_length=1)


class StartOnboardingRequest(BaseModel):
    patient: PatientContext
    category: CategoryContext
    language: str = "en"


class StartOnboardingResponse(BaseModel):
    success: bool
    session_id: str
    patient_id: str
    category: CategoryContext
    language: str
    status: str
    stage: str
    message: str


class SendMessageRequest(BaseModel):
    message: str = Field(..., min_length=1)


class SendMessageResponse(BaseModel):
    success: bool
    session_id: str
    stage: str
    message: str
    status: str


class AIStateUpdate(BaseModel):
    reason_for_visit: Optional[Any] = None
    symptoms: Optional[Any] = None
    medical_history: Optional[Any] = None
    medications: Optional[Any] = None
    allergies: Optional[Any] = None


class AIOnboardingResponse(BaseModel):
    reply: str
    next_stage: str
    state_updates: AIStateUpdate
    completed: bool = False


class Symptom(BaseModel):
    name: str
    duration: Optional[str] = None
    severity: Optional[str] = None
    description: Optional[str] = None


class PatientOnboardingResult(BaseModel):
    patient_id: str
    abha_id: Optional[str] = None
    phone: Optional[str] = None
    insurance_scheme: Optional[str] = None
    category: dict[str, Any]
    chief_complaint: Optional[str] = None
    symptoms: list[Symptom] = Field(default_factory=list)
    medical_history: list[str] = Field(default_factory=list)
    medications: list[str] = Field(default_factory=list)
    allergies: list[str] = Field(default_factory=list)
    dietary_habits: Optional[str] = None
    additional_information: Optional[str] = None
    language: str
    onboarding_completed: bool = True



class CompleteOnboardingRequest(BaseModel):
    confirmed: bool = False


class CompleteOnboardingResponse(BaseModel):
    success: bool
    session_id: str
    status: str
    result: Optional[PatientOnboardingResult] = None
    message: str