"""
Rewrite node for LangGraph.

This node takes the reflection produced by the ReflectNode and rewrites
it to a more formal style. The output is a dictionary containing
the key 'rewritten'.
"""

from langgraph.graph import node
from typing import Dict, Any


class RewriteNode:
    """
    A LangGraph node that rewrites the reflection message.
    """

    @node
    def run(self, state: Dict[str, Any]) -> Dict[str, str]:
        """
        Rewrite the reflection message.

        Parameters
        ----------
        state : dict
            The current state of the graph. Expected to contain a 'reflection'
            key with the message produced by the ReflectNode.

        Returns
        -------
        dict
            A dictionary with a single key 'rewritten' containing the
            rewritten message.
        """
        reflection = state.get("reflection", "")
        # Simple rewrite: replace "I see" with "I notice"
        rewritten = reflection.replace("I see", "I notice")
        return {"rewritten": rewritten}