from typing_extensions import TypedDict
from datetime import date
from typing import Annotated, Dict, Any
from langgraph.graph.message import add_messages

class PainLevels(TypedDict):
    left_head: int
    right_head: int
    left_arm: int
    left_hand: int
    right_arm: int
    right_hand: int
    left_body_trunk: int
    right_body_trunk: int
    left_leg: int
    left_foot: int
    right_leg: int
    right_foot: int

class Surgery(TypedDict):
    surgery_name: str
    time: date

class PatientID(TypedDict):
    name: str
    DOB: date
    gender: str
    contact: str
    emergency_contact: str

class MainSymptom(TypedDict):
    main_symptom: str
    symptom_length: str

class Pain(TypedDict):
    pain_location: str
    pain_side: str
    pain_intensity: int
    pain_description: str
    start_time: date
    radiation: bool
    triggers: str
    symptom: str

class MedicalHistory(TypedDict):
    medical_condition: str
    first_time: date
    surgery_history: str
    medication: str
    allergy: str

class FamilyHistory(TypedDict):
    family_history: str

class SocialHistory(TypedDict):
    occupation: str
    smoke: bool
    alcohol: bool
    drug: bool
    support_system: str
    living_condition: str

class ReviewSystem(TypedDict):
    weight_change: str
    fever: bool
    chill: bool
    night_sweats: bool
    sleep: str
    gastrointestinal: str
    urinary: str

class PainManagement(TypedDict):
    pain_medication: str
    specialist: bool
    other_therapy: str
    effectiveness: bool

class Functional(TypedDict):
    life_quality: str
    limit_activity: str
    mood: str

class Plan(TypedDict):
    goal: str
    expectation: str
    alternative_treatment_illness: str

class PatientData(TypedDict):
    ID: PatientID
    symptom: MainSymptom
    pain: Pain
    medical_hist: MedicalHistory
    family_hist: FamilyHistory
    social_hist: SocialHistory
    review_system: ReviewSystem
    pain_manage: PainManagement
    functional: Functional
    plan: Plan

class DataState(TypedDict):
    """State representing the patient's data status and conversation."""
    messages: Annotated[list, add_messages]
    data: Dict[str, PatientData]
    finished: bool