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


class PatientContext(BaseModel):
    id: str = Field(..., min_length=1)
    name: Optional[str] = None
    age: Optional[int] = Field(None, ge=0, le=120)
    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, max_length=5000)


class ErrorDetail(BaseModel):
    code: str
    message: str


class ErrorResponse(BaseModel):
    error: ErrorDetail


class SendMessageResponse(BaseModel):
    success: bool
    session_id: str
    stage: str
    message: str
    status: str
    response_time_ms: Optional[float] = None
    response_time_sec: Optional[float] = None


class Symptom(BaseModel):
    name: str
    duration: Optional[str] = None
    severity: Optional[str] = None
    description: Optional[str] = None
    location: Optional[str] = None
    onset: Optional[str] = None
    frequency: Optional[str] = None
    pattern: Optional[str] = None
    triggers: Optional[str] = None
    relieving_factors: Optional[str] = None
    aggravating_factors: Optional[str] = None


class MedicationDetail(BaseModel):
    name: str
    dosage: Optional[str] = None
    frequency: Optional[str] = None
    duration: Optional[str] = None
    reason: Optional[str] = None


class AllergyDetail(BaseModel):
    name: str
    type: Optional[str] = None
    reaction: Optional[str] = None


class MedicalHistoryDetail(BaseModel):
    condition: str
    duration: Optional[str] = None
    status: Optional[str] = None


class SurgicalHistoryDetail(BaseModel):
    procedure: str
    approximate_date: Optional[str] = None
    reason: Optional[str] = None


class PreviousEpisodeDetail(BaseModel):
    description: str
    time: Optional[str] = None
    outcome: Optional[str] = None


class FamilyHistoryDetail(BaseModel):
    condition: str
    relationship: Optional[str] = None


class AdditionalInfoDetail(BaseModel):
    category: str
    value: 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
    surgical_history: Optional[Any] = None
    previous_episodes: Optional[Any] = None
    family_history: Optional[Any] = None
    additional_information: Optional[Any] = None


class AIOnboardingResponse(BaseModel):
    reply: str
    next_stage: str
    state_updates: AIStateUpdate
    completed: bool = False
    suggested_specialty: Optional[dict[str, Any]] = 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[Union[Symptom, dict, str]] = Field(default_factory=list)
    medical_history: list[Union[str, dict, MedicalHistoryDetail]] = Field(default_factory=list)
    medications: list[Union[str, dict, MedicationDetail]] = Field(default_factory=list)
    allergies: list[Union[str, dict, AllergyDetail]] = Field(default_factory=list)
    surgical_history: Optional[list[Union[str, dict, SurgicalHistoryDetail]]] = Field(default_factory=list)
    previous_episodes: Optional[list[Union[str, dict, PreviousEpisodeDetail]]] = Field(default_factory=list)
    family_history: Optional[list[Union[str, dict, FamilyHistoryDetail]]] = Field(default_factory=list)
    dietary_habits: Optional[str] = None
    additional_information: Optional[Any] = 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
    response_time_ms: Optional[float] = None
    response_time_sec: Optional[float] = None