"""
Agent Graph - LangGraph State Machine for Document Navigation

Implements the Navigator Agent with full RLM (Recursive Language Model) support:

1. Decomposes queries into sub-questions (Section 2.2 - Recursive Loop)
2. Navigates the document tree (expand/traverse decisions)
3. Stores findings as pointers (Variable Stitching - Section 2.2)
4. Synthesizes final answer from stored pointers
5. Tree of Thoughts prompting (Section 7.2)
6. Recursive sub-LLM invocation for complex queries

Agent State follows Appendix C specification:
- question: Current question being answered
- sub_questions: Decomposed sub-questions (via LLM)
- current_node_id: Where we are in the document tree
- visited_nodes: Navigation history
- variables: Stored findings as $POINTER -> content
- pending_questions: Sub-questions not yet answered
- context: Accumulated context (pointers only!)
- answer: Final synthesized answer
- trace: Full retrieval trace for transparency
"""

from __future__ import annotations

import re
from datetime import datetime, timezone
from typing import Any, Literal, TypedDict, cast

import structlog

from rnsr.agent.variable_store import VariableStore, generate_pointer_name
from rnsr.indexing.kv_store import KVStore
from rnsr.indexing.semantic_search import SemanticSearcher
from rnsr.models import SkeletonNode, TraceEntry

logger = structlog.get_logger(__name__)


# =============================================================================
# Agent State Definition (Appendix C)
# =============================================================================


class AgentState(TypedDict):
    """
    State for the RNSR Navigator Agent.
    
    All fields use pointer-based Variable Stitching:
    - Full content stored in VariableStore
    - Only $POINTER names in state fields
    """
    
    # Query processing
    question: str
    sub_questions: list[str]
    pending_questions: list[str]
    current_sub_question: str | None
    
    # Navigation state
    current_node_id: str | None
    visited_nodes: list[str]
    navigation_path: list[str]
    
    # Tree of Thoughts (ToT) state - Section 7.2
    nodes_to_visit: list[str] # Queue for parallel exploration
    scored_candidates: list[dict[str, Any]]  # [{node_id, score, reasoning}]
    backtrack_stack: list[str]  # Stack of parent node IDs for backtracking
    dead_ends: list[str]  # Nodes marked as dead ends
    top_k: int  # Number of top candidates to explore
    tot_selection_threshold: float  # Minimum probability for selection
    tot_dead_end_threshold: float   # Probability threshold for dead end
    
    # Variable stitching (pointers only!)
    variables: list[str]  # List of $POINTER names
    context: str  # Contains pointers, not full content
    
    # Output
    answer: str | None
    confidence: float
    
    # Question metadata (e.g., multiple-choice options)
    metadata: dict[str, Any]
    
    # Traceability
    trace: list[dict[str, Any]]
    iteration: int
    max_iterations: int


# =============================================================================
# Navigator Tools
# =============================================================================


def create_navigator_tools(
    skeleton: dict[str, SkeletonNode],
    kv_store: KVStore,
    variable_store: VariableStore,
) -> dict[str, Any]:
    """
    Create tools for the Navigator Agent.
    
    Args:
        skeleton: Skeleton index (node_id -> SkeletonNode).
        kv_store: KV store with full content.
        variable_store: Variable store for findings.
        
    Returns:
        Dictionary of tool functions.
    """
    
    def get_node_summary(node_id: str) -> str:
        """Get the summary of a skeleton node."""
        node = skeleton.get(node_id)
        if node is None:
            return f"Node {node_id} not found"
        
        children_info = ""
        if node.child_ids:
            children = [skeleton.get(cid) for cid in node.child_ids]
            children_info = "\nChildren:\n" + "\n".join(
                f"  - {c.node_id}: {c.header}" for c in children if c
            )
        
        return f"""
Node: {node.node_id}
Header: {node.header}
Level: {node.level}
Summary: {node.summary}
{children_info}
"""
    
    def navigate_to_child(node_id: str, child_id: str) -> str:
        """Navigate to a child node."""
        node = skeleton.get(node_id)
        if node is None:
            return f"Node {node_id} not found"
        
        if child_id not in node.child_ids:
            return f"{child_id} is not a child of {node_id}"
        
        child = skeleton.get(child_id)
        if child is None:
            return f"Child {child_id} not found"
        
        return get_node_summary(child_id)
    
    def expand_node(node_id: str) -> str:
        """
        EXPAND: Fetch full content and store as variable.
        
        Use when the summary answers the question.
        """
        node = skeleton.get(node_id)
        if node is None:
            return f"Node {node_id} not found"
        
        # Fetch full content from KV store
        content = kv_store.get(node_id)
        if content is None:
            return f"No content found for {node_id}"
        
        # Generate pointer name
        pointer = generate_pointer_name(node.header)
        
        # Store as variable
        variable_store.assign(pointer, content, node_id)
        
        logger.info(
            "node_expanded",
            node_id=node_id,
            pointer=pointer,
            chars=len(content),
        )
        
        return f"Stored content as {pointer} ({len(content)} chars)"
    
    def store_finding(
        pointer_name: str,
        content: str,
        source_node_id: str = "",
    ) -> str:
        """
        Store a finding as a pointer variable.
        
        Args:
            pointer_name: Name like $LIABILITY_CLAUSE (must start with $)
            content: Full text content to store
            source_node_id: Source node for traceability
        """
        if not pointer_name.startswith("$"):
            pointer_name = "$" + pointer_name.upper()
        
        try:
            meta = variable_store.assign(pointer_name, content, source_node_id)
            return f"Stored as {pointer_name} ({meta.char_count} chars)"
        except Exception as e:
            return f"Error storing: {e}"
    
    def compare_variables(*pointers: str) -> str:
        """
        Compare multiple stored variables.
        
        Resolves pointers and returns content for comparison.
        Use during synthesis phase only.
        """
        results = []
        for pointer in pointers:
            content = variable_store.resolve(pointer)
            if content:
                results.append(f"=== {pointer} ===\n{content}")
            else:
                results.append(f"=== {pointer} ===\n[Not found]")
        
        return "\n\n".join(results)
    
    def list_stored_variables() -> str:
        """List all stored variables with metadata."""
        variables = variable_store.list_variables()
        if not variables:
            return "No variables stored yet."
        
        lines = ["Stored Variables:"]
        for v in variables:
            lines.append(f"  {v.pointer}: {v.char_count} chars from {v.source_node_id}")
        
        return "\n".join(lines)
    
    def synthesize_from_variables() -> str:
        """
        Get all stored content for final synthesis.
        
        Call this at the end to get full content for answer generation.
        """
        variables = variable_store.list_variables()
        if not variables:
            return "No variables to synthesize from."
        
        parts = []
        for v in variables:
            content = variable_store.resolve(v.pointer)
            if content:
                parts.append(f"=== {v.pointer} ===\n{content}")
        
        return "\n\n".join(parts)
    
    return {
        "get_node_summary": get_node_summary,
        "navigate_to_child": navigate_to_child,
        "expand_node": expand_node,
        "store_finding": store_finding,
        "compare_variables": compare_variables,
        "list_stored_variables": list_stored_variables,
        "synthesize_from_variables": synthesize_from_variables,
    }


# =============================================================================
# State Management Functions
# =============================================================================


def create_initial_state(
    question: str,
    root_node_id: str,
    max_iterations: int = 20,
    top_k: int = 3,
    metadata: dict[str, Any] | None = None,
    tot_selection_threshold: float = 0.4,
    tot_dead_end_threshold: float = 0.1,
) -> AgentState:
    """Create the initial agent state with ToT support."""
    return AgentState(
        question=question,
        sub_questions=[],
        pending_questions=[],
        current_sub_question=None,
        current_node_id=root_node_id,
        visited_nodes=[],
        navigation_path=[root_node_id],
        # Tree of Thoughts state
        nodes_to_visit=[],
        scored_candidates=[],
        backtrack_stack=[],
        dead_ends=[],
        top_k=top_k,
        tot_selection_threshold=tot_selection_threshold,
        tot_dead_end_threshold=tot_dead_end_threshold,
        # Variable stitching
        variables=[],
        context="",
        answer=None,
        confidence=0.0,
        metadata=metadata or {},
        trace=[],
        iteration=0,
        max_iterations=max_iterations,
    )


def add_trace_entry(
    state: AgentState,
    node_type: Literal["decomposition", "navigation", "variable_stitching", "synthesis"],
    action: str,
    details: dict | None = None,
) -> None:
    """Add a trace entry to the state."""
    entry = TraceEntry(
        timestamp=datetime.now(timezone.utc).isoformat(),
        node_type=node_type,
        action=action,
        details=details or {},
    )
    state["trace"].append(entry.model_dump())


# =============================================================================
# Tree of Thoughts (ToT) Prompting - Section 7.2
# =============================================================================

# ToT System Prompt as specified in the research paper
TOT_SYSTEM_PROMPT = """You are a Deep Research Agent navigating a document tree.

You are currently at Node: {current_node_summary}
Children Nodes: {children_summaries}
Your Goal: {query}

EVALUATION TASK:
For each child node, estimate the probability (0.0 to 1.0) that it contains relevant information for the goal.

INSTRUCTIONS:
1. Evaluate: For each child node, analyze its summary and estimate relevance probability.
2. Be OPEN-MINDED: Select nodes with probability > {selection_threshold} (moderate evidence of relevance).
3. Look for matches: Prefer nodes with facts/entities mentioned in the query, but also consider broad thematic matches.
4. Balance PRECISION and RECALL: Do not prune branches too early. If unsure, include the node.
5. Plan: Select the top-{top_k} most promising nodes.
6. Reasoning: Explain briefly what SPECIFIC content in the summary makes it relevant.
7. Backtrack Signal: If NO child seems relevant (all probabilities < {dead_end_threshold}), report "DEAD_END".

OUTPUT FORMAT (JSON):
{{
    "evaluations": [
        {{"node_id": "...", "probability": 0.85, "reasoning": "Summary mentions X which directly relates to query about Y"}},
        {{"node_id": "...", "probability": 0.60, "reasoning": "Contains information about Z"}},
        ...
    ],
    "selected_nodes": ["node_id_1", "node_id_2", ...],
    "is_dead_end": false,
    "backtrack_reason": null
}}

If this is a dead end:
{{
    "evaluations": [...],
    "selected_nodes": [],
    "is_dead_end": true,
    "backtrack_reason": "None of the children appear to contain information about X."
}}

Respond ONLY with the JSON, no other text."""


def _format_node_summary(node: SkeletonNode) -> str:
    """Format a node's summary for the ToT prompt."""
    return f"[{node.node_id}] {node.header}: {node.summary or '(no summary)'}"


def _format_children_summaries(
    skeleton: dict[str, SkeletonNode],
    child_ids: list[str],
) -> str:
    """Format all children summaries for the ToT prompt."""
    if not child_ids:
        return "(no children - this is a leaf node)"
    
    lines = []
    for child_id in child_ids:
        child = skeleton.get(child_id)
        if child:
            lines.append(f"  - {_format_node_summary(child)}")
    
    return "\n".join(lines) if lines else "(no children found)"


def evaluate_children_with_tot(
    state: AgentState,
    skeleton: dict[str, SkeletonNode],
    top_k_override: int | None = None,
) -> dict[str, Any]:
    """
    Use Tree of Thoughts prompting to evaluate child nodes.
    
    This implements Section 7.2 of the research paper:
    "We use Tree of Thoughts (ToT) which explicitly encourages the model 
    to explore multiple branches of the index before committing to a path."
    
    Args:
        state: Current agent state.
        skeleton: Skeleton index.
        top_k_override: Optional override for adaptive exploration.
        
    Returns:
        Dictionary with evaluations, selected_nodes, is_dead_end, and backtrack_reason.
    """
    node_id = state.get("current_node_id")
    if node_id is None:
        return {"evaluations": [], "selected_nodes": [], "is_dead_end": True, "backtrack_reason": "No current node"}
    
    node = skeleton.get(node_id)
    if node is None:
        return {"evaluations": [], "selected_nodes": [], "is_dead_end": True, "backtrack_reason": "Node not found"}
    
    # If no children, this is a leaf - expand instead of traverse
    if not node.child_ids:
        return {"evaluations": [], "selected_nodes": [], "is_dead_end": False, "backtrack_reason": None, "is_leaf": True}
    
    # Format the ToT prompt
    current_summary = _format_node_summary(node)
    children_summaries = _format_children_summaries(skeleton, node.child_ids)
    query = state.get("current_sub_question") or state.get("question", "")
    top_k = top_k_override if top_k_override is not None else state.get("top_k", 3)
    selection_threshold = state.get("tot_selection_threshold", 0.4)
    dead_end_threshold = state.get("tot_dead_end_threshold", 0.1)
    
    prompt = TOT_SYSTEM_PROMPT.format(
        current_node_summary=current_summary,
        children_summaries=children_summaries,
        query=query,
        top_k=top_k,
        selection_threshold=selection_threshold,
        dead_end_threshold=dead_end_threshold,
    )
    
    # Call the LLM
    try:
        from rnsr.llm import get_llm
        import json
        
        llm = get_llm()
        response = llm.complete(prompt)
        response_text = str(response).strip()
        
        # Parse JSON response with robust error handling
        try:
            # Handle potential markdown code blocks
            if "```" in response_text:
                match = re.search(r'\{[\s\S]*\}', response_text)
                if match:
                    response_text = match.group(0)

            result = json.loads(response_text)
        except json.JSONDecodeError:
            logger.warning("tot_json_repair_attempt", original_response=response_text)
            # Fallback to asking the LLM to fix the JSON
            repair_prompt = f"""The following text is a malformed JSON object. Please fix it and return ONLY the corrected JSON. Do not add any commentary.

Malformed JSON:
{response_text}"""
            from rnsr.llm import get_llm

            llm = get_llm()
            repaired_response_text = str(llm.complete(repair_prompt)).strip()
            
            # Final attempt to parse the repaired JSON
            if "```" in repaired_response_text:
                 match = re.search(r'\{[\s\S]*\}', repaired_response_text)
                 if match:
                    repaired_response_text = match.group(0)

            result = json.loads(repaired_response_text)
        
        logger.debug(
            "tot_evaluation_complete",
            node_id=node_id,
            evaluations=len(result.get("evaluations", [])),
            selected=result.get("selected_nodes", []),
            is_dead_end=result.get("is_dead_end", False),
        )
        
        return result
        
    except Exception as e:
        logger.warning("tot_evaluation_failed", error=str(e), node_id=node_id)
        
        # Fallback: use simple heuristic (select first unvisited children)
        visited = state.get("visited_nodes", [])
        dead_ends = state.get("dead_ends", [])
        
        unvisited = [
            cid for cid in node.child_ids 
            if cid not in visited and cid not in dead_ends
        ]
        
        return {
            "evaluations": [{"node_id": cid, "probability": 0.5, "reasoning": "Fallback selection"} for cid in unvisited[:top_k]],
            "selected_nodes": unvisited[:top_k],
            "is_dead_end": len(unvisited) == 0,
            "backtrack_reason": "No unvisited children" if len(unvisited) == 0 else None,
        }


def backtrack_to_parent(
    state: AgentState,
    skeleton: dict[str, SkeletonNode],
) -> AgentState:
    """
    Backtrack to parent node when current path is a dead end.
    
    Implements the backtracking logic from Section 7.2:
    "Backtrack: If a node yields no useful info, report 'Dead End' and return to parent."
    """
    new_state = cast(AgentState, dict(state))
    
    current_id = new_state["current_node_id"]
    
    # Mark current node as dead end
    if current_id and current_id not in new_state["dead_ends"]:
        new_state["dead_ends"].append(current_id)
    
    # Try to backtrack using the backtrack stack
    if new_state["backtrack_stack"]:
        parent_id = new_state["backtrack_stack"].pop()
        new_state["current_node_id"] = parent_id
        new_state["navigation_path"].append(parent_id)
        
        add_trace_entry(
            new_state,
            "navigation",
            f"Backtracked to {parent_id} (dead end at {current_id})",
            {"from": current_id, "to": parent_id, "reason": "dead_end"},
        )
        
        logger.debug("backtrack_success", from_node=current_id, to_node=parent_id)
    else:
        # No parent to backtrack to - we're done exploring
        add_trace_entry(
            new_state,
            "navigation",
            "Cannot backtrack - at root or stack empty",
            {"current": current_id},
        )
        logger.debug("backtrack_failed", reason="empty_stack")
    
    new_state["iteration"] = new_state["iteration"] + 1
    return new_state


# =============================================================================
# LangGraph Node Functions
# =============================================================================


# Decomposition prompt for RLM recursive sub-task generation
DECOMPOSITION_PROMPT = """You are analyzing a complex query to decompose it into sub-tasks.

Query: {query}

Document Structure (top-level sections):
{structure}

TASK: Decompose this query into specific sub-tasks that can be executed independently.

RULES:
1. Each sub-task should target a specific section or piece of information
2. Sub-tasks should be answerable by reading specific document sections
3. For comparison queries, create one sub-task per item being compared
4. For multi-hop queries, create sequential sub-tasks (find A, then use A to find B)
5. Maximum 5 sub-tasks to maintain efficiency

OUTPUT FORMAT (JSON):
{{
    "sub_tasks": [
        {{"id": 1, "task": "Find X in section Y", "target_section": "section_hint"}},
        {{"id": 2, "task": "Extract Z from W", "target_section": "section_hint"}}
    ],
    "synthesis_plan": "How to combine sub-task results into final answer"
}}

Return ONLY valid JSON."""


def decompose_query(state: AgentState) -> AgentState:
    """
    Decompose the main question into sub-questions using LLM.
    
    Implements Section 2.2 "The Recursive Loop":
    "The model's capability to divide a complex reasoning task into 
    smaller, manageable sub-tasks and invoke instances of itself to solve them."
    """
    new_state = cast(AgentState, dict(state))  # Copy
    question = new_state["question"]
    
    # Try LLM-based decomposition
    try:
        from rnsr.llm import get_llm
        import json
        
        llm = get_llm()
        
        # Get document structure for context
        # Note: skeleton is not available here, so we'll do basic decomposition
        structure_hint = "Document sections available for navigation"
        
        prompt = DECOMPOSITION_PROMPT.format(
            query=question,
            structure=structure_hint,
        )
        
        response = llm.complete(prompt)
        response_text = str(response).strip()
        
        # Extract JSON from response
        json_match = re.search(r'\{[\s\S]*\}', response_text)
        if json_match:
            decomposition = json.loads(json_match.group())
            sub_tasks = decomposition.get("sub_tasks", [])
            
            if sub_tasks:
                sub_questions = [t["task"] for t in sub_tasks]
                new_state["sub_questions"] = sub_questions
                new_state["pending_questions"] = sub_questions.copy()
                new_state["current_sub_question"] = sub_questions[0]
                
                add_trace_entry(
                    new_state,
                    "decomposition",
                    f"LLM decomposed into {len(sub_questions)} sub-tasks",
                    {
                        "sub_questions": sub_questions,
                        "synthesis_plan": decomposition.get("synthesis_plan", ""),
                    },
                )
                
                new_state["iteration"] = new_state["iteration"] + 1
                return new_state
    
    except Exception as e:
        logger.warning("llm_decomposition_failed", error=str(e))
    
    # Fallback: simple decomposition patterns
    sub_questions = _simple_decompose(question)
    new_state["sub_questions"] = sub_questions
    new_state["pending_questions"] = sub_questions.copy()
    new_state["current_sub_question"] = sub_questions[0] if sub_questions else question
    
    add_trace_entry(
        new_state,
        "decomposition",
        f"Decomposed into {len(sub_questions)} sub-questions (fallback)",
        {"sub_questions": sub_questions},
    )
    
    new_state["iteration"] = new_state["iteration"] + 1
    return new_state


def _simple_decompose(question: str) -> list[str]:
    """
    Simple pattern-based decomposition fallback.
    
    Handles common query patterns without LLM.
    """
    question_lower = question.lower()
    
    # Pattern: "Compare X and Y" or "X vs Y"
    compare_patterns = [
        r"compare\s+(.+?)\s+(?:and|with|to|vs\.?)\s+(.+)",
        r"difference(?:s)?\s+between\s+(.+?)\s+and\s+(.+)",
        r"(.+?)\s+vs\.?\s+(.+)",
    ]
    
    for pattern in compare_patterns:
        match = re.search(pattern, question_lower)
        if match:
            item1, item2 = match.groups()
            return [
                f"Find information about {item1.strip()}",
                f"Find information about {item2.strip()}",
            ]
    
    # Pattern: "What are the X in sections A, B, and C"
    multi_section = re.search(
        r"in\s+(?:sections?\s+)?(.+?,\s*.+?(?:,\s*.+)*)",
        question_lower,
    )
    if multi_section:
        sections = [s.strip() for s in multi_section.group(1).split(",")]
        return [f"Find relevant information in {s}" for s in sections[:5]]
    
    # Pattern: "List all X" or "Find all X" - may need iteration
    if re.search(r"(list|find|show)\s+all", question_lower):
        return [
            f"Search for all relevant sections",
            question,  # Original as synthesis query
        ]
    
    # Default: single question
    return [question]


def navigate_tree(
    state: AgentState,
    skeleton: dict[str, SkeletonNode],
) -> AgentState:
    """
    Navigate the document tree based on current sub-question.
    
    This function now handles the node visitation queue for multi-path exploration.
    """
    new_state = cast(AgentState, dict(state))

    # Add detailed logging to debug navigation loops
    logger.debug(
        "navigate_step_start",
        iteration=new_state["iteration"],
        current_node=new_state["current_node_id"],
        queue=new_state["nodes_to_visit"],
        visited=new_state["visited_nodes"],
    )

    # If current node is None, try to pop from the visit queue
    if new_state["current_node_id"] is None and new_state["nodes_to_visit"]:
        next_node_id = new_state["nodes_to_visit"].pop(0)
        new_state["current_node_id"] = next_node_id
        
        # Also add to navigation path
        if next_node_id not in new_state["navigation_path"]:
            new_state["navigation_path"].append(next_node_id)
            
        add_trace_entry(
            new_state,
            "navigation",
            f"Visiting next node from queue: {next_node_id}",
            {"queue_size": len(new_state["nodes_to_visit"])},
        )
    
    node_id = new_state["current_node_id"]
    if node_id is None:
        add_trace_entry(new_state, "navigation", "No current node to navigate to and queue is empty")
        return new_state
    
    node = skeleton.get(node_id)
    
    if node is None:
        add_trace_entry(new_state, "navigation", f"Node {node_id} not found")
        return new_state
    
    # Don't add to visited here - let expand/traverse handle that
    # to avoid preventing expansion of newly queued nodes
    
    add_trace_entry(
        new_state,
        "navigation",
        f"At node {node_id}: {node.header}",
        {"summary": node.summary, "children": node.child_ids},
    )
    
    new_state["iteration"] = new_state["iteration"] + 1
    return new_state


def should_expand(
    state: AgentState,
    skeleton: dict[str, SkeletonNode],
) -> Literal["expand", "traverse", "backtrack", "done"]:
    """
    Decide whether to expand (fetch content), traverse (go to children), 
    backtrack (dead end), or finish.
    
    Uses Tree of Thoughts evaluation to make intelligent decisions.
    """
    node_id = state.get("current_node_id")
    if node_id is None:
        return "done"
    
    node = skeleton.get(node_id)
    if node is None:
        return "done"
    
    # Check iteration limit
    if state.get("iteration", 0) >= state.get("max_iterations", 20):
        return "done"
    
    # Removed variable count limit - iteration limit is sufficient
    
    # If no children, must expand (leaf node)
    # But only if not already visited
    visited = state.get("visited_nodes", [])
    if not node.child_ids:
        if node_id in visited:
            # Already expanded this leaf, done with it
            return "done"
        return "expand"
    
    # Check if all children are visited or dead ends
    visited = state.get("visited_nodes", [])
    dead_ends = state.get("dead_ends", [])
    unvisited_children = [
        cid for cid in node.child_ids 
        if cid not in visited and cid not in dead_ends
    ]
    
    if not unvisited_children:
        # All children explored - backtrack or done
        if state.get("backtrack_stack"):
            return "backtrack"
        return "done"
    
    # Adaptive exploration: increase top_k if we haven't found much information yet
    base_top_k = state.get("top_k", 3)
    variables_found = len(state.get("variables", []))
    iteration = state.get("iteration", 0)
    
    if iteration > 5 and variables_found == 0:
        adaptive_top_k = min(len(node.child_ids), base_top_k * 3)
    elif iteration > 3 and variables_found < 2:
        adaptive_top_k = min(len(node.child_ids), base_top_k * 2)
    else:
        adaptive_top_k = base_top_k

    # Use ToT evaluation to decide
    tot_result = evaluate_children_with_tot(state, skeleton, top_k_override=adaptive_top_k)
    
    # Check for dead end signal from ToT
    if tot_result.get("is_dead_end", False):
        if state.get("backtrack_stack"):
            return "backtrack"
        return "done"
    
    # Check if ToT says this is a leaf (should expand)
    if tot_result.get("is_leaf", False):
        return "expand"
    
    # If ToT selected nodes, traverse
    if tot_result.get("selected_nodes"):
        return "traverse"
    
    # Fallback: expand if at deep level (H3+)
    if node.level >= 3:
        return "expand"
    
    # Otherwise traverse to children
    return "traverse"


def expand_current_node(
    state: AgentState,
    skeleton: dict[str, SkeletonNode],
    kv_store: KVStore,
    variable_store: VariableStore,
) -> AgentState:
    """
    EXPAND: Fetch full content and store as variable.
    """
    new_state = cast(AgentState, dict(state))
    
    node_id = new_state["current_node_id"]
    if node_id is None:
        return new_state
    
    # Check if already visited to prevent infinite loops
    if node_id in new_state.get("visited_nodes", []):
        add_trace_entry(
            new_state,
            "navigation", 
            f"Skipping already-visited node {node_id}",
            {"node_id": node_id},
        )
        new_state["current_node_id"] = None
        new_state["iteration"] = new_state["iteration"] + 1
        return new_state
    
    node = skeleton.get(node_id)
    
    if node is None:
        new_state["current_node_id"] = None
        new_state["iteration"] = new_state["iteration"] + 1
        return new_state
    
    # Fetch full content
    content = kv_store.get(node_id)
    if content is None:
        add_trace_entry(new_state, "variable_stitching", f"No content for {node_id}")
        new_state["current_node_id"] = None
        new_state["iteration"] = new_state["iteration"] + 1
        return new_state
    
    # Generate and store as variable
    pointer = generate_pointer_name(node.header)
    variable_store.assign(pointer, content, node_id)
    
    new_state["variables"].append(pointer)
    new_state["context"] += f"\nFound: {pointer} (from {node.header})"
    
    # Mark node as visited and clear current node
    if node_id not in new_state["visited_nodes"]:
        new_state["visited_nodes"].append(node_id)
    new_state["current_node_id"] = None  # Process queue next
    
    add_trace_entry(
        new_state,
        "variable_stitching",
        f"Stored {pointer}",
        {"node_id": node_id, "header": node.header, "chars": len(content)},
    )
    
    new_state["iteration"] = new_state["iteration"] + 1
    return new_state


def traverse_to_children(
    state: AgentState,
    skeleton: dict[str, SkeletonNode],
    semantic_searcher: SemanticSearcher | None = None,
) -> AgentState:
    """
    TRAVERSE: Navigate to child nodes using Tree of Thoughts reasoning.
    
    Primary Strategy: Tree of Thoughts (ToT) with adaptive exploration.
    - Uses LLM reasoning to evaluate child nodes based on summaries
    - Dynamically increases exploration when insufficient variables found
    - Preserves document structure and context (per research paper Section 7.2)
    
    Optional: Semantic search as shortcut for finding entry points (Section 9.1)
    """
    new_state = cast(AgentState, dict(state))
    
    node_id = new_state["current_node_id"]
    if node_id is None:
        return new_state
    
    node = skeleton.get(node_id)
    
    if node is None or not node.child_ids:
        return new_state
    
    question = new_state["question"]
    base_top_k = new_state.get("top_k", 3)
    
    # Adaptive exploration: increase top_k if we haven't found much information yet
    # This addresses the original issue: "agent should be able to explore all nodes if it needs to"
    variables_found = len(new_state.get("variables", []))
    iteration = new_state.get("iteration", 0)
    
    # Dynamic top_k based on progress
    if iteration > 5 and variables_found == 0:
        # Not finding anything - expand search significantly
        adaptive_top_k = min(len(node.child_ids), base_top_k * 3)
        add_trace_entry(
            new_state,
            "navigation",
            f"Expanding search: {variables_found} variables after {iteration} iterations",
            {"base_top_k": base_top_k, "adaptive_top_k": adaptive_top_k},
        )
    elif iteration > 3 and variables_found < 2:
        # Finding very little - expand moderately
        adaptive_top_k = min(len(node.child_ids), base_top_k * 2)
    else:
        # Normal exploration
        adaptive_top_k = base_top_k
    
    # Use Tree of Thoughts as primary navigation method (per research paper)
    tot_result = evaluate_children_with_tot(new_state, skeleton, top_k_override=adaptive_top_k)
    
    # Check for dead end
    if tot_result.get("is_dead_end", False):
        add_trace_entry(
            new_state,
            "navigation",
            f"Dead end detected at {node_id}: {tot_result.get('backtrack_reason', 'Unknown')}",
            {"node_id": node_id, "backtrack_reason": tot_result.get("backtrack_reason")},
        )
        # Mark as dead end and don't change current node
        if node_id not in new_state["dead_ends"]:
            new_state["dead_ends"].append(node_id)
        new_state["iteration"] = new_state["iteration"] + 1
        return new_state
    
    selected_nodes = tot_result.get("selected_nodes", [])
    evaluations = tot_result.get("evaluations", [])
    new_state["scored_candidates"] = evaluations

    # Queue selected children for visitation
    if selected_nodes:
        new_state["nodes_to_visit"].extend(selected_nodes)
        # Mark parent as visited since we've traversed its children
        if node_id not in new_state["visited_nodes"]:
            new_state["visited_nodes"].append(node_id)
        # Unset current node to force the next navigate step to pull from the queue
        new_state["current_node_id"] = None
        add_trace_entry(
            new_state,
            "navigation",
            f"Queued {len(selected_nodes)} children from {node_id}",
            {"nodes": selected_nodes, "parent": node_id},
        )
    else:
        # No candidates available - this should trigger backtracking
        add_trace_entry(
            new_state,
            "navigation",
            f"No unvisited candidates at {node_id}",
            {"node_id": node_id},
        )
    
    new_state["iteration"] = new_state["iteration"] + 1
    return new_state


def synthesize_answer(
    state: AgentState,
    variable_store: VariableStore,
) -> AgentState:
    """
    Synthesize final answer from stored variables using an LLM.
    
    Uses the configured LLM to generate a concise answer
    from the resolved variable content.
    """
    new_state = cast(AgentState, dict(state))
    
    # Resolve all variables
    pointers = new_state["variables"]
    
    if not pointers:
        new_state["answer"] = "No relevant content found."
        new_state["confidence"] = 0.0
    else:
        # Collect all content
        contents = []
        for pointer in pointers:
            content = variable_store.resolve(pointer)
            if content:
                contents.append(content)
        
        context_text = "\n\n---\n\n".join(contents)
        question = new_state["question"]
        
        # Use LLM to synthesize answer
        try:
            from rnsr.llm import get_llm

            llm = get_llm()

            metadata = new_state.get("metadata", {})
            options = metadata.get("options")

            if options and isinstance(options, list) and len(options) > 0:
                # QuALITY multiple-choice question
                # Ensure each option is a string and format with letters
                options_text = "\n".join([f"{chr(65+i)}. {str(opt)}" for i, opt in enumerate(options)])
                prompt = f"""Based on the provided context, answer this multiple-choice question.

Question: {question}

Options:
{options_text}

Context:
{context_text}

Instructions:
1. Read the context carefully
2. Determine which option is best supported by the evidence in the context
3. Respond with ONLY the letter and full text of the correct option
4. Format: "X. [complete option text]" where X is A, B, C, or D

Your answer:"""
            else:
                # Standard open-ended question
                prompt = f"""Based on the following context, answer the question concisely.

IMPORTANT INSTRUCTIONS:
- If the context contains information that DIRECTLY answers the question, provide that answer
- If the context contains information that allows you to INFER the answer, provide your best inference
- Use evidence from the context to support your answer
- Only say "Cannot determine from available context" if the context is completely unrelated or missing critical information
- It's better to provide a reasonable answer based on available evidence than to say "cannot determine"

Question: {question}

Context:
{context_text}

Answer (be concise, direct, and confident):"""

            response = llm.complete(prompt)
            answer = str(response).strip()
            
            # Normalize multiple-choice answers
            if options:
                answer_lower = answer.lower().strip()
                matched = False
                
                for i, opt in enumerate(options):
                    letter = chr(65 + i)  # A, B, C, D
                    opt_lower = opt.lower()
                    
                    # Match patterns: "A.", "A)", "(A)", "A. answer text"
                    if (answer_lower.startswith(f"{letter.lower()}.") or 
                        answer_lower.startswith(f"{letter.lower()})") or 
                        answer_lower.startswith(f"({letter.lower()})")):
                        answer = opt
                        matched = True
                        break
                    
                    # Exact match (case insensitive)
                    if answer_lower == opt_lower:
                        answer = opt
                        matched = True
                        break
                    
                    # Check if option text is contained in answer
                    if opt_lower in answer_lower or answer_lower in opt_lower:
                        # Use similarity to pick best match if multiple partial matches
                        answer = opt
                        matched = True
                        break
                
                new_state["confidence"] = 0.8 if matched else 0.1
            else:
                new_state["confidence"] = min(1.0, len(pointers) * 0.3)
            
            new_state["answer"] = answer

        except Exception as e:
            logger.warning("llm_synthesis_failed", error=str(e))
            # Fallback to context concatenation
            new_state["answer"] = f"Context found:\n{context_text}"
            new_state["confidence"] = min(1.0, len(pointers) * 0.2)
    
    add_trace_entry(
        new_state,
        "synthesis",
        "Generated answer from variables",
        {"variables_used": pointers, "confidence": new_state["confidence"]},
    )
    
    return new_state


# =============================================================================
# RLM Recursive Execution Functions (Section 2.2)
# =============================================================================


def execute_sub_task_with_llm(
    sub_task: str,
    context: str,
    llm_fn: Any = None,
) -> str:
    """
    Execute a single sub-task using an LLM.
    
    Implements the recursive LLM pattern from Section 2.2:
    "For each contract, call the LLM API (invoking itself) with a 
    specific sub-prompt: 'Extract liability clause from this text.'"
    
    Args:
        sub_task: The sub-task description/prompt.
        context: The context to process.
        llm_fn: Optional LLM function. If None, uses default.
        
    Returns:
        LLM response as string.
    """
    if llm_fn is None:
        try:
            from rnsr.llm import get_llm
            llm = get_llm()
            llm_fn = lambda p: str(llm.complete(p))
        except Exception as e:
            logger.warning("llm_not_available", error=str(e))
            return f"[Error: LLM not available - {str(e)}]"
    
    prompt = f"""{sub_task}

Context:
{context}

Response:"""
    
    try:
        return llm_fn(prompt)
    except Exception as e:
        logger.error("sub_task_execution_failed", error=str(e))
        return f"[Error: {str(e)}]"


def batch_execute_sub_tasks(
    sub_tasks: list[str],
    contexts: list[str],
    batch_size: int = 5,
    max_parallel: int = 4,
) -> list[str]:
    """
    Execute multiple sub-tasks in parallel batches.
    
    Implements Section 2.3 "Optimization via Batching":
    "Instead of making 1,000 individual calls to summarize 1,000 paragraphs,
    the RLM writes code to group paragraphs into chunks of 5 and processes 
    them in parallel threads."
    
    Args:
        sub_tasks: List of sub-task prompts.
        contexts: List of contexts for each sub-task.
        batch_size: Items per batch (default 5).
        max_parallel: Max parallel threads (default 4).
        
    Returns:
        List of results for each sub-task.
    """
    from concurrent.futures import ThreadPoolExecutor, as_completed
    
    if len(sub_tasks) != len(contexts):
        raise ValueError("sub_tasks and contexts must have same length")
    
    if not sub_tasks:
        return []
    
    logger.info(
        "batch_execution_start",
        num_tasks=len(sub_tasks),
        batch_size=batch_size,
    )
    
    # Get LLM function
    try:
        from rnsr.llm import get_llm
        llm = get_llm()
        llm_fn = lambda p: str(llm.complete(p))
    except Exception as e:
        logger.error("batch_llm_failed", error=str(e))
        return [f"[Error: {str(e)}]"] * len(sub_tasks)
    
    results: list[str] = [""] * len(sub_tasks)
    
    with ThreadPoolExecutor(max_workers=max_parallel) as executor:
        future_to_idx: dict[Any, int] = {}
        
        for idx, (task, ctx) in enumerate(zip(sub_tasks, contexts)):
            future = executor.submit(execute_sub_task_with_llm, task, ctx, llm_fn)
            future_to_idx[future] = idx
        
        for future in as_completed(future_to_idx):
            idx = future_to_idx[future]
            try:
                results[idx] = future.result(timeout=120)
            except Exception as e:
                results[idx] = f"[Error: {str(e)}]"
    
    logger.info("batch_execution_complete", num_results=len(results))
    return results


def process_pending_questions(
    state: AgentState,
    skeleton: dict[str, SkeletonNode],
    kv_store: KVStore,
    variable_store: VariableStore,
) -> AgentState:
    """
    Process all pending sub-questions using recursive LLM calls.
    
    This implements the full RLM recursive loop:
    1. For each pending question, find relevant context
    2. Invoke sub-LLM to extract answer
    3. Store result as variable
    4. Move to next question
    """
    new_state = cast(AgentState, dict(state))
    pending = new_state.get("pending_questions", [])
    
    if not pending:
        return new_state
    
    current_question = pending[0]
    
    # Find relevant content for this question
    # Use current node and its children
    node_id = new_state.get("current_node_id") or "root"
    node = skeleton.get(node_id)
    
    if node is None:
        # Pop and continue
        new_state["pending_questions"] = pending[1:]
        return new_state
    
    # Get context from current node and children
    context_parts = []
    
    # Add current node content
    current_content = kv_store.get(node_id) if node_id else None
    if current_content:
        context_parts.append(f"[{node.header}]\n{current_content}")
    
    # Add children summaries
    for child_id in node.child_ids[:5]:  # Limit to 5 children
        child = skeleton.get(child_id)
        if child:
            child_content = kv_store.get(child_id)
            if child_content:
                context_parts.append(f"[{child.header}]\n{child_content[:2000]}")
    
    context = "\n\n---\n\n".join(context_parts)
    
    # Execute sub-task with LLM
    result = execute_sub_task_with_llm(
        sub_task=f"Answer this question: {current_question}",
        context=context,
    )
    
    # Store as variable
    pointer = generate_pointer_name(current_question[:30])
    variable_store.assign(pointer, result, node_id or "root")
    new_state["variables"].append(pointer)
    
    add_trace_entry(
        new_state,
        "decomposition",
        f"Processed sub-question via recursive LLM",
        {
            "question": current_question,
            "pointer": pointer,
            "context_length": len(context),
        },
    )
    
    # Pop processed question
    new_state["pending_questions"] = pending[1:]
    if pending[1:]:
        new_state["current_sub_question"] = pending[1]
    
    new_state["iteration"] = new_state["iteration"] + 1
    return new_state


# =============================================================================
# Graph Builder (LangGraph)
# =============================================================================


def build_navigator_graph(
    skeleton: dict[str, SkeletonNode],
    kv_store: KVStore,
    semantic_searcher: SemanticSearcher | None = None,
) -> Any:
    """
    Build the LangGraph state machine for document navigation.
    
    Implements Tree of Thoughts (ToT) navigation with backtracking support.
    
    Returns a compiled graph that can be invoked with a question.
    
    Usage:
        graph = build_navigator_graph(skeleton, kv_store)
        result = graph.invoke({"question": "What are the payment terms?"})
    """
    try:
        from langgraph.graph import END, StateGraph
    except ImportError:
        raise ImportError(
            "LangGraph not installed. Install with: pip install langgraph"
        )
    
    variable_store = VariableStore()
    
    # Create graph
    graph = StateGraph(AgentState)
    
    # Add nodes
    graph.add_node("decompose", decompose_query)
    
    graph.add_node(
        "navigate",
        lambda state: navigate_tree(cast(AgentState, state), skeleton),
    )
    
    graph.add_node(
        "expand",
        lambda state: expand_current_node(cast(AgentState, state), skeleton, kv_store, variable_store),
    )
    
    graph.add_node(
        "traverse",
        lambda state: traverse_to_children(
            cast(AgentState, state),
            skeleton,
            semantic_searcher,
        ),
    )
    
    # Add backtrack node for ToT dead-end handling
    graph.add_node(
        "backtrack",
        lambda state: backtrack_to_parent(cast(AgentState, state), skeleton),
    )
    
    graph.add_node(
        "synthesize",
        lambda state: synthesize_answer(cast(AgentState, state), variable_store),
    )
    
    # Add edges
    graph.add_edge("decompose", "navigate")
    
    # Conditional edge based on expand/traverse/backtrack decision
    graph.add_conditional_edges(
        "navigate",
        lambda s: should_expand(s, skeleton),
        {
            "expand": "expand",
            "traverse": "traverse",
            "backtrack": "backtrack",
            "done": "synthesize",
        },
    )
    
    # After expand, traverse, or backtrack, always return to the main navigation
    # handler, which will decide what to do next (e.g., pull from queue)
    graph.add_edge("expand", "navigate")
    graph.add_edge("traverse", "navigate")
    graph.add_edge("backtrack", "navigate")
    
    graph.add_edge("synthesize", END)
    
    # Set entry point
    graph.set_entry_point("decompose")
    
    logger.info("navigator_graph_built", features=["ToT", "backtracking"])
    
    return graph.compile()


# =============================================================================
# High-Level API
# =============================================================================


def run_navigator(
    question: str,
    skeleton: dict[str, SkeletonNode],
    kv_store: KVStore,
    max_iterations: int = 20,
    top_k: int | None = None,
    use_semantic_search: bool = True,
    semantic_searcher: SemanticSearcher | None = None,
    metadata: dict[str, Any] | None = None,
    tot_selection_threshold: float = 0.4,
    tot_dead_end_threshold: float = 0.1,
) -> dict[str, Any]:
    """
    Run the navigator agent on a question.
    
    Args:
        question: User's question.
        skeleton: Skeleton index.
        kv_store: KV store with full content.
        max_iterations: Maximum navigation iterations.
        top_k: Number of top children to explore (default: auto-detect based on tree depth).
        use_semantic_search: Use semantic search (O(log N)) instead of ToT evaluation (O(N)).
            Allows exploring ALL leaf nodes ranked by relevance, preventing missed data.
        semantic_searcher: Optional pre-built semantic searcher. If None and use_semantic_search=True, creates one.
        tot_selection_threshold: Minimum probability for ToT node selection (0.0-1.0).
        tot_dead_end_threshold: Probability threshold for declaring a dead end (0.0-1.0).
        
    Returns:
        Dictionary with answer, confidence, trace.
        
    Example:
        result = run_navigator(
            "What are the liability terms?",
            skeleton,
            kv_store,
            use_semantic_search=True,  # Enable semantic search
        )
        print(result["answer"])
    """
    # Get root node
    root_id = None
    root_node = None
    for node in skeleton.values():
        if node.level == 0:
            root_id = node.node_id
            root_node = node
            break
    
    if root_id is None:
        return {
            "answer": "Error: No root node found in skeleton index.",
            "confidence": 0.0,
            "trace": [],
        }
    
    # Auto-detect top_k based on tree structure
    if top_k is None:
        num_root_children = len(root_node.child_ids) if root_node else 0
        if num_root_children > 10:
            # Flat hierarchy (e.g., QuALITY): explore more children
            top_k = min(10, num_root_children)
        else:
            # Deep hierarchy (e.g., PDFs): explore fewer
            top_k = 3
    
    # Semantic search disabled by default per research paper Section 9.1:
    # "Hybridize Search: Give the agent a vector_search tool as a SHORTCUT.
    # The agent can use vector search to find a starting node and then
    # SWITCH TO TREE TRAVERSAL for local exploration."
    # 
    # Primary navigation uses ToT reasoning-based retrieval.
    if use_semantic_search and semantic_searcher is None:
        try:
            from rnsr.indexing.semantic_search import create_semantic_searcher
            semantic_searcher = create_semantic_searcher(skeleton, kv_store)
            logger.info(
                "semantic_search_optional_tool",
                nodes=len(skeleton),
                note="Available as shortcut for entry points only",
            )
        except Exception as e:
            logger.warning(
                "semantic_search_unavailable",
                error=str(e),
            )
            semantic_searcher = None
    
    logger.info(
        "using_tot_reasoning_navigation",
        method="Tree of Thoughts",
        adaptive_exploration=True,
        note="LLM reasons about document structure to navigate",
    )
    
    # Build and run graph
    graph = build_navigator_graph(skeleton, kv_store, semantic_searcher)
    
    initial_state = create_initial_state(
        question=question,
        root_node_id=root_id,
        max_iterations=max_iterations,
        top_k=top_k,
        metadata=metadata,
        tot_selection_threshold=tot_selection_threshold,
        tot_dead_end_threshold=tot_dead_end_threshold,
    )
    
    final_state = graph.invoke(initial_state)
    
    return {
        "answer": final_state.get("answer", ""),
        "confidence": final_state.get("confidence", 0.0),
        "trace": final_state.get("trace", []),
        "variables_used": final_state.get("variables", []),
        "nodes_visited": final_state.get("visited_nodes", []),
    }
