# sutram_services/agent_state.py

from typing import List, Union, Optional, Any, Dict, TYPE_CHECKING

from langchain_core.messages import BaseMessage
from pydantic import BaseModel, Field

if TYPE_CHECKING:
    pass


class PlanStep(BaseModel):
    """A single step in the agent's execution plan."""
    tool_name: str = Field(..., description="The exact name of the tool to execute")
    step_name: str = Field(..., description="A unique identifier for this plan step")
    tool_args: Dict[str, Any] = Field(default_factory=dict, description="Arguments for the tool call")
    dependencies: List[str] = Field(default_factory=list, description="List of step names this step depends on")
    status: str = Field(default="pending", description="Current status of the step")

class MacroPlanResponse(BaseModel):
    """Response model for the LLM's planning output."""
    plan: List[PlanStep] = Field(default_factory=list, description="List of plan steps")
    reflection: Optional[str] = Field(default="", description="Reasoning behind the plan")

# Define the Agent's state structure for LangGraph
class AgentState(BaseModel):
    """
    Represents the state of the LangGraph agent execution for Macro Planning.
    """
    # Core Graph Invocation
    input: str = Field(default="") # Use Field for optional/default values
    trace_id: str = Field(default="") # Use Field for optional/default values
    otel_trace_id: str = Field(default="") # Use Field for optional/default values
    span_id: Optional[str] = None

    user_id: Optional[str] = None,
    metadata: Optional[Dict[str, Any]] = None,

    # 🎯 NEW ACE FIELDS
    active_template_id: Optional[str] = None
    context_playbook: Optional[str] = None

    # LangChain Chat History and Scratchpad
    chat_history: List[BaseMessage] = Field(default_factory=list)
    agent_scratchpad: List[Union[BaseMessage, Any]] = Field(default_factory=list) # For LLM's thinking/internal monologue
    # RAG/Context fields (Used before/during planning)
    retrieved_context: Optional[str] = None
    macro_plan: List[Dict[str, Any]] = Field(default_factory=list)  # This should match the PlanStep structure
    context_relevance_level: Optional[str] = Field(
        default=None,
        description="Relevance level of the retrieved context (e.g., 'high', 'medium', 'low', 'error')"
    )
    status: Optional[str] = Field(
        default="in_progress",
        description="Current Agent State (e.g., 'in_progress', 'completed', 'failed')"
    )
    critique_data: Optional[Any] = Field(default=None, description="Reflection critique data")

    # Evaluation fields (Used at the end)
    context_relevance_score: Optional[float] = None
    answer_relevance_score: Optional[float] = None
    groundedness_score: Optional[float] = None

    relevant_tool_names: List[str] = Field(default_factory=list,
                                           description="List of tool names relevant to the current query.")
    relevant_tools: List[Any] = Field(default_factory=list,
                                    description="List of tool objects relevant to the current query.")
    relevant_tools_description: str = Field(default="",
                                            description="formatted descriptions of relevant tools for the LLM prompt.")

    # Flag for internal routing
    next_step: str  = None # e.g., 'call_tool', 'generate_answer', 'check_groundedness', 'end'

    # Overall error flag (Early termination)
    error: Optional[str] = None

    # -----------------------------------------------------------
    # PHASE 3: MACRO PLANNING FIELDS (Core Orchestration Data)
    # -----------------------------------------------------------

    # 1. The Plan (Generated/Refined by LLM)
    # [{"tool_name": str, "tool_input": dict, "status": "pending"|..., "depends_on": List[str]}]
    macro_plan: Optional[List[Dict[str, Any]]] = Field(default_factory=list)
    evaluation: Optional[Dict[str, Any]] = Field(default_factory=dict)

    # 2. The Execution Results (Durable storage from Temporal runs)
    # {"step_name": {"tool_name": str, "result": Any, "error": Optional[str], "temporal_run_id": str}}
    macro_context: Dict[str, Any] = Field(default_factory=dict)

    # 3. The Feedback (Input to the next planning step)
    reflection_data: Dict[str, Any] = Field(default_factory=dict)

    # 4. Final Output (Read by chat_endpoint)
    final_response_content: Optional[str]  =  None # CONSOLIDATED DEFINITION

    # Use the safe .get() method in your router without crashing
    def get(self, key: str, default: Any = None) -> Any:
        return getattr(self, key, default)