from typing import TypedDict, Annotated, Optional, List, Dict, Any
from langgraph.graph.message import add_messages

class AgentState(TypedDict, total=False):
    # Conversation
    messages: Annotated[list, add_messages]
    user_input: str

    # Routing
    intent: Optional[str]
    intent_confidence: Optional[float]
    sub_intents: Optional[List[str]]

    # Entity Context (multi-turn memory)
    last_order_id: Optional[str]
    last_product_id: Optional[str]
    last_seller_id: Optional[str]
    last_category: Optional[str]
    session_context: Dict[str, Any]

    # Tool Outputs
    sql_query: Optional[str]
    sql_result: Optional[Dict[str, Any]]
    rag_result: Optional[str]
    rag_score: Optional[float]

    # Sentiment
    frustration_score: Optional[float]
    emotions_detected: Optional[List[str]]

    # Proactive Logic
    is_late_delivery: Optional[bool]
    days_overdue: Optional[float]
    compensation_offered: Optional[bool]
    compensation_tier: Optional[str]

    # Escalation
    escalation_required: Optional[bool]
    escalation_reason: Optional[str]
    case_id: Optional[str]
    escalation_summary: Optional[Dict[str, Any]]

    # Control Flow
    retry_count: int
    error_log: List[str]
    sanitized: Optional[bool]
    final_response: Optional[str]

    # Clarification
    clarification_needed: Optional[bool]
    clarification_type: Optional[str]


def init_state(user_input: str) -> AgentState:
    """Initialize state with safe defaults to prevent KeyError"""
    return {
        "user_input": user_input,
        "messages": [],
        "session_context": {},
        "error_log": [],
        "retry_count": 0,
        "intent": None,
        "intent_confidence": None,
        "sub_intents": None,
        "last_order_id": None,
        "last_product_id": None,
        "last_seller_id": None,
        "last_category": None,
        "sql_query": None,
        "sql_result": None,
        "rag_result": None,
        "rag_score": None,
        "frustration_score": None,
        "emotions_detected": None,
        "is_late_delivery": None,
        "days_overdue": None,
        "compensation_offered": None,
        "compensation_tier": None,
        "escalation_required": None,
        "escalation_reason": None,
        "case_id": None,
        "sanitized": None,
        "final_response": None,
        "clarification_needed": None,
        "clarification_type": None,
    }