"""
Variable extraction and resolution for ResolutionFlow tree structures.

Supports three variable tokens:
- [USER_INPUT:prompt] — prompts user for input during session navigation
- [VAR:name] — references a previously saved variable value
- [SAVE_AS:name] — saves the current context as a named variable
"""
import re
from dataclasses import dataclass
from typing import Any


@dataclass
class VariableDefinition:
    """A variable found in a tree structure."""
    name: str
    kind: str  # 'user_input', 'reference', 'save_as'
    prompt: str  # For user_input: the prompt text; for others: empty
    node_id: str  # The node where this variable appears


# Regex patterns for variable tokens
USER_INPUT_RE = re.compile(r'\[USER_INPUT:([^\]]+)\]')
VAR_REF_RE = re.compile(r'\[VAR:([^\]]+)\]')
SAVE_AS_RE = re.compile(r'\[SAVE_AS:([^\]]+)\]')


def extract_variables(tree_structure: dict[str, Any]) -> list[VariableDefinition]:
    """Extract all variable definitions from a tree structure.

    Traverses the tree recursively and finds all [USER_INPUT:...],
    [VAR:...], and [SAVE_AS:...] tokens.

    Args:
        tree_structure: The recursive tree_structure dict.

    Returns:
        List of VariableDefinition objects found in the tree.
    """
    variables: list[VariableDefinition] = []
    _scan_node(tree_structure, variables)
    return variables


def _scan_node(node: dict[str, Any], variables: list[VariableDefinition]) -> None:
    """Scan a node and its children for variable tokens."""
    node_id = node.get("id", "unknown")

    # Collect all text fields to scan
    text_fields = [
        node.get("question", ""),
        node.get("title", ""),
        node.get("description", ""),
        node.get("help_text", ""),
        node.get("expected_outcome", ""),
        node.get("solution", ""),
    ]

    # Include option labels
    for opt in node.get("options", []):
        text_fields.append(opt.get("label", ""))

    # Include resolution steps
    for step in node.get("resolution_steps", []):
        text_fields.append(step)

    # Include commands
    for cmd in node.get("commands", []):
        text_fields.append(cmd)

    # Scan all text fields
    for text in text_fields:
        if not text:
            continue

        for match in USER_INPUT_RE.finditer(text):
            variables.append(VariableDefinition(
                name=match.group(1).strip(),
                kind="user_input",
                prompt=match.group(1).strip(),
                node_id=node_id,
            ))

        for match in VAR_REF_RE.finditer(text):
            variables.append(VariableDefinition(
                name=match.group(1).strip(),
                kind="reference",
                prompt="",
                node_id=node_id,
            ))

        for match in SAVE_AS_RE.finditer(text):
            variables.append(VariableDefinition(
                name=match.group(1).strip(),
                kind="save_as",
                prompt="",
                node_id=node_id,
            ))

    # Recurse into children
    for child in node.get("children", []):
        _scan_node(child, variables)


def resolve_variables(text: str, variables: dict[str, str]) -> str:
    """Replace variable tokens in text with their resolved values.

    - [VAR:name] → replaced with variables[name] if present
    - [USER_INPUT:prompt] → replaced with variables[prompt] if present
    - [SAVE_AS:name] → removed (save directives don't appear in output)

    Args:
        text: The text containing variable tokens.
        variables: Dict mapping variable names to their values.

    Returns:
        Text with variable tokens replaced.
    """
    def _replace_var(match: re.Match) -> str:
        name = match.group(1).strip()
        return variables.get(name, f"[VAR:{name}]")

    def _replace_input(match: re.Match) -> str:
        prompt = match.group(1).strip()
        return variables.get(prompt, f"[USER_INPUT:{prompt}]")

    def _replace_save(match: re.Match) -> str:
        return ""  # Save directives are removed from output

    result = VAR_REF_RE.sub(_replace_var, text)
    result = USER_INPUT_RE.sub(_replace_input, result)
    result = SAVE_AS_RE.sub(_replace_save, result)
    return result