"""
state.py

Shared LangGraph state definitions for WalkXR. This file should stay generic 
enough to support future agents, not just the Small Moments roleplay agent.
"""


from __future__ import annotations
from typing import Any, Dict, List, Literal, Optional, TypedDict


# Named shortcuts for common string literals
# ------------------------------------------
Role = Literal["system", "user", "assistant", "tool"]
Phase = Literal[
    "input",
    "memory_loaded",
    "context_retrieved",
    "prompt_constructed",
    "llm_called",
    "response_generated",
    "history_updated",
    "memory_updated",
    "completed",
    "error",
]

# Safe response in case LLM does not generate a user-friendly response
SAFE_RESPONSE = "Thank you for sharing your thoughts. Let me take a moment to process the feelings you have shared with me. Are there additional details you would like to discuss?"

# Message-level state
# -------------------
class MessageRecord(TypedDict):
    """
    Message format used inside the graph state.

    We use a dict-based structure instead of raw tuples so it can scale more
    easily to:
    - System messages
    - Tool messages
    - Timestamps / metadata
    - Future multi-agent handoffs

    Fields:
    - role:
        Who produced the message (user, assistant, system, or tool).

    - content:
        The text content of the message.
    """
    role: Role
    content: str


# Retrieval state
# ---------------
class RetrievedChunk(TypedDict, total=False):
    """
    One retrieved context item from the RAG pipeline.

    total=False means some keys are optional. This is useful because the
    current retrieval code mainly returns text strings, but later the retrieval
    engine may expose other metadata such as score, source, or document id.

    Fields:
    - text:
        The actual content retrieved from the knowledgebase.

    - source:
        Identifier of where this chunk came from.
        Example: "small_moments_doc_v1".

    - score:
        Relevance score from the retrieval system (higher = more relevant).

    - metadata:
        Additional structured info (section, tags, etc.).
        Useful for debugging or advanced filtering.
    """
    text: str
    source: str
    score: float
    metadata: Dict[str, Any]


# Memory state
# ------------
class UserMemory(TypedDict, total=False):
    """
    Memory about the user or current walk session.

    Keep this minimal for EPIC 3. Later EPICs can expand to cross-session memory.

    Fields:
    - summary:
        A rolling summary of the conversation so far.
        Used to maintain continuity without passing the full transcript.

    - preferences:
        Notable user likes, dislikes, or style preferences.
        Example: "prefers reflective questions", "doesn't like small talk".

    - facts:
        Concrete facts about the user.
        Example: "starting a new job", "lives alone", "has social anxiety".

    - open_loops:
        Unresolved topics or threads that the agent may want to revisit.
        Example: "user mentioned fear of networking but didn't explore it".

    - prior_insights:
        Meaningful reflections or realizations the user has already had.
        Helps avoid repeating the same insights.
    """
    summary: str
    preferences: List[str]
    facts: List[str]
    open_loops: List[str]
    prior_insights: List[str]


# Constructed prompt state
# ------------------------
class PromptContext(TypedDict, total=False):
    """
    Structured payload prepared for the LLM node.

    This is not the final rendered prompt string. Instead, it is a packaged
    subset of graph state that the LLM node will format into the final prompt.

    Fields:
    - user_input:
        The current user message for this turn.

    - recent_messages:
        A recent slice of the conversation history for local continuity.

    - memory:
        Memory about the user or current session.

    - retrieved_context:
        External context retrieved for the current turn.

    - current_step:
        The current step in the conversation/workflow.

    - feedback:
        The feedback provided by the supervisor to try again.
    """
    user_input: str
    recent_messages: List[MessageRecord]
    memory: UserMemory
    retrieved_context: List[RetrievedChunk]
    current_step: str
    feedback: str   # provides brief justification for supervisor score


# Emotional state placeholder
# ---------------------------
class EmotionalState(TypedDict, total=False):
    """
    Placeholder for the emotional state engine described in the roadmap.

    These fields are optional because the emotional analysis system does not
    need to be fully implemented yet, but the state should make room for it.

    Purpose: Give the system a structured way to track the user's emotional 
    state, which can inform the agent's responses and reflections.
    """
    valence: float
    arousal: float
    openness: float
    friction: float
    primary_emotion: str
    tone: str


# Agent output state
# ------------------
class AgentOutput(TypedDict, total=False):
    """
    Structured output from an agent node.

    Fields:
    - response_text:
        The final text response shown to the user.

    - source_chunks:
        Text snippets used to ground the response.

    - should_end_turn:
        Whether the agent considers this interaction complete.

    - next_step:
        Suggested next step in the workflow.
        Example: "reflection", "follow_up", "end".

    - handoff_to:
        Name of another agent to transfer control to (future multi-agent use).
    """
    response_text: str
    source_chunks: List[RetrievedChunk]
    should_end_turn: bool
    next_step: str
    handoff_to: str


# Top-level graph state
# ---------------------
class WalkState(TypedDict, total=False):
    """
    Shared state for the WalkXR LangGraph workflow.

    This is the main state object the graph will read from and write to.

    Sections:
    - Identity / session:
        Tracks who the user is and where they are in the experience.

    - Interactions:
        The current input, constructed prompt, and conversation history.

    - Retrieval:
        Context pulled from external knowledge (RAG).

    - Memory:
        Current knowledge / understanding of the user.

    - Agent output:
        The result of the most recent agent step.

    - Control / orchestration:
        Internal graph execution state (phase, errors, scratchpad).

    - Evaluation:
        Debugging, tracing, and scoring information.

    - Testing:
        Used solely by test scripts.
    """

    # Identity / session
    user_id: str
    session_id: str
    walk_id: str
    active_agent: str
    current_step: str
    turn_index: int

    # Interactions
    user_input: str
    conversation_history: List[MessageRecord]

    # Retrieval
    retrieval_query: str
    retrieved_context: List[RetrievedChunk]

    # Memory
    memory: UserMemory

    # Prompt constructions
    prompt_context: PromptContext

    # Future state engines
    emotional_state: EmotionalState

    # Agent result
    agent_output: AgentOutput

    # Control / orchestration
    phase: Phase                        # Current phase of the graph execution
    agent_scratchpad: Dict[str, Any]    # Temporary workspace for agents to store intermediate thoughts, tool calls, etc.
    error: Optional[str]                # Error message if something went wrong, else None

    # Evaluation
    supervisor_score: str               # Scores current agent_output as "good" or "bad"
    supervisor_retries: int             # Count number of retries to prevent exceeding max limit
    max_retries_allowed: int            # Maximum limit for number of retries allowed
    safe_response: str                  # Fail-safe response if maximum retries reached
    debug_print: bool

    # Testing
    bad_response_test: Optional[str]    # Replaces LLM response with custom response to test self-reflection loop

# Helper constructors
# -------------------
def create_initial_state(
    *,
    user_input: str,
    user_id: Optional[str] = None,
    session_id: Optional[str] = None,
    walk_id: Optional[str] = None,
    active_agent: str = "small_moments_roleplay",
    current_step: str = "conversation",
    turn_index: int = 0,
    conversation_history: Optional[List[MessageRecord]] = None,
    max_retries_allowed: int = 3,
    safe_response: str = SAFE_RESPONSE,
    bad_response_test: Optional[str] = None,
    debug_print: bool = False
) -> WalkState:
    """
    Creates an initial WalkState used to start a new graph.
    """
    return WalkState(
        user_id=user_id or "unknown_user",
        session_id=session_id or "default_session",
        walk_id=walk_id or "default_walk",
        active_agent=active_agent,
        current_step=current_step,
        turn_index=turn_index,
        user_input=user_input,
        conversation_history=conversation_history or [],
        retrieval_query=user_input,
        retrieved_context=[],
        memory=UserMemory(
            summary="",
            preferences=[],
            facts=[],
            open_loops=[],
            prior_insights=[],
        ),
        prompt_context=PromptContext(
            user_input="",
            recent_messages=[],
            memory=UserMemory(
                summary="",
                preferences=[],
                facts=[],
                open_loops=[],
                prior_insights=[],
            ),
            retrieved_context=[],
            current_step=current_step,
            feedback="",
        ),
        emotional_state=EmotionalState(),
        agent_output=AgentOutput(
            response_text="",
            source_chunks=[],
            should_end_turn=False,
        ),
        phase="input",
        agent_scratchpad={},
        error=None,
        supervisor_score="good",
        supervisor_retries=0,
        max_retries_allowed=max_retries_allowed,
        safe_response=safe_response,
        debug_print=debug_print,
        bad_response_test=bad_response_test
    )


def append_message(state: WalkState, *, role: Role, content: str) -> WalkState:
    """
    Convenience helper for adding a message to state.

    Instead of needing to append the message list manually in every node, this 
    helper centralizes the logic and keeps the nodes.py code cleaner.
    """
    state.setdefault("conversation_history", [])
    state["conversation_history"].append(MessageRecord(role=role, content=content))
    return state