"""
VSS State Schema
TypedDict for LangGraph orchestration state.
"""

from typing import TypedDict, Literal, Optional, List, Dict, Any
from datetime import datetime, timezone
import uuid


class VSSState(TypedDict, total=False):
    """Unified state for Mode 1 query workflow."""
    
    # Request identification
    request_id: str
    timestamp: str
    status: Literal["pending", "routing", "processing", "completed", "failed"]
    
    # User input
    user_query: str
    uploaded_media: Optional[Dict[str, Any]]  # {"type": "image", "data": bytes}
    
    # Routing
    intent: Literal["sql", "vlm", "unknown"]
    intent_confidence: float
    routing_reason: str
    
    # SQL processing
    sql_query: Optional[str]
    sql_result: Optional[List[Dict]]
    sql_result_formatted: Optional[str]
    sql_error: Optional[str]
    sql_retry_count: int
    
    # VLM processing
    vlm_prompt: Optional[str]
    vlm_response: Optional[str]
    vlm_error: Optional[str]
    
    # Escalation
    escalated: bool
    escalation_reason: Optional[str]
    
    # Final output
    final_response: Optional[str]
    response_source: Literal["sql", "vlm", "error", ""]
    
    # Tracing
    trace: List[str]
    error: Optional[str]
    processing_time_ms: Optional[int]


def create_state(query: str, media: Optional[Dict] = None) -> VSSState:
    """Create initial state for a query."""
    return VSSState(
        request_id=str(uuid.uuid4()),
        timestamp=datetime.now(timezone.utc).isoformat(),
        status="pending",
        user_query=query,
        uploaded_media=media,
        intent="unknown",
        intent_confidence=0.0,
        routing_reason="",
        sql_query=None,
        sql_result=None,
        sql_result_formatted=None,
        sql_error=None,
        sql_retry_count=0,
        vlm_prompt=None,
        vlm_response=None,
        vlm_error=None,
        escalated=False,
        escalation_reason=None,
        final_response=None,
        response_source="",
        trace=[f"State initialized: {query[:100]}"],
        error=None,
        processing_time_ms=None,
    )


def add_trace(state: VSSState, msg: str) -> None:
    """Add trace message."""
    state["trace"].append(f"[{datetime.now(timezone.utc).strftime('%H:%M:%S')}] {msg}")


def set_error(state: VSSState, error: str) -> None:
    """Set error state."""
    state["error"] = error
    state["status"] = "failed"
    add_trace(state, f"ERROR: {error}")