"""
LangGraph State Definition for Secure Reasoning Agent
Tracks all state through the Chain-of-Checks workflow.
"""
from typing import TypedDict, List, Optional, Annotated
from operator import add
from langchain_core.messages import BaseMessage

from schemas import (
    ExecutionPlan,
    SafetyCheckResult,
    CryptoLogEntry,
    Justification,
    ExecutionResult
)


class AgentState(TypedDict):
    """
    State tracked throughout the LangGraph execution.
    
    The state flows through: Plan → Safety Check → Execute → Log → Justify → Loop
    """
    
    # ========================================================================
    # CONVERSATION & CONTEXT
    # ========================================================================
    messages: Annotated[List[BaseMessage], add]
    """Chat history with the user and internal LLM calls. Uses 'add' reducer to append."""
    
    task: str
    """The original user task/query."""
    
    task_id: str
    """Unique identifier for this execution (for audit trail)."""
    
    user_id: Optional[str]
    """Optional user identifier for multi-user environments."""
    
    # ========================================================================
    # PLANNING STATE
    # ========================================================================
    plan: Optional[ExecutionPlan]
    """The generated execution plan with all steps."""
    
    current_step_index: int
    """Which step we're currently processing (0-indexed)."""
    
    # ========================================================================
    # SAFETY & VALIDATION
    # ========================================================================
    safety_status: Optional[SafetyCheckResult]
    """Result of safety check for current step. None if not yet checked."""
    
    safety_blocked: bool
    """Quick flag: True if any step was blocked by safety guardrails."""
    
    # ========================================================================
    # EXECUTION STATE
    # ========================================================================
    execution_result: Optional[ExecutionResult]
    """Result from executing the current step."""
    
    final_result: Optional[str]
    """The final answer/output when all steps complete."""
    
    # ========================================================================
    # AUDIT TRAIL & CRYPTOGRAPHIC LOGGING
    # ========================================================================
    logs: List[CryptoLogEntry]
    """Cryptographic proofs for each executed step (Merkle roots, hashes, etc.)."""
    
    justifications: List[Justification]
    """Agent's reasoning for each action taken."""
    
    # ========================================================================
    # ERROR HANDLING
    # ========================================================================
    error: Optional[str]
    """Error message if execution fails."""
    
    status: str
    """Current execution status: 'planning', 'executing', 'completed', 'failed', 'blocked'."""


def create_initial_state(task: str, task_id: str, user_id: Optional[str] = None) -> AgentState:
    """
    Factory function to create a fresh AgentState for a new task.
    
    Args:
        task: The user's task/query
        task_id: Unique identifier for this execution
        user_id: Optional user identifier
        
    Returns:
        Initialized AgentState ready for LangGraph processing
    """
    return AgentState(
        messages=[],
        task=task,
        task_id=task_id,
        user_id=user_id,
        plan=None,
        current_step_index=0,
        safety_status=None,
        safety_blocked=False,
        execution_result=None,
        final_result=None,
        logs=[],
        justifications=[],
        error=None,
        status="planning"
    )