from typing import TypedDict, Annotated, List, Union
from langchain_core.messages import BaseMessage
import operator

class AgentState(TypedDict):
    # Standard message history
    messages: Annotated[List[BaseMessage], operator.add]
    
    # Metadata for routing
    user_role: str  # patient, clinician, researcher
    intent_type: str  # diagnosis, treatment, monitoring, general
    trace_id: str
    
    # CDM / FHIR Data
    session_id: str
    patient_id: str
    fhir_data: List[dict]
    trend_analysis: str
    trend_data: dict
    current_medications: List[dict]
    
    # Validation status
    is_valid: bool
    is_safe: bool
    
    # Intermediate outputs
    patient_response: str
    clinician_outputs: List[str]
    research_output: str
    
    # Recovery loop counter
    attempts: int
    
    # Sources for information
    sources: Annotated[List[str], operator.add]
    
    # Evidence citations for clinical recommendations (Bug 12.3)
    evidence_citations: Annotated[List[dict], operator.add]
    
    # Execution logs
    logs: Annotated[List[str], operator.add]
    
    # Metrics (tokens, time)
    metrics: Annotated[List[dict], operator.add]