from typing import Tuple

from langchain_experimental.tot.memory import ToTDFSMemory
from langchain_experimental.tot.thought import ThoughtValidity


class ToTController:
    """
    Tree of Thought (ToT) controller.

    This is a version of a ToT controller, dubbed in the paper as a "Simple
    Controller".

    It has one parameter `c` which is the number of children to explore for each
    thought.
    """

    def __init__(self, c: int = 3):
        """
        Initialize the controller.

        Args:
            c: The number of children to explore at each node.
        """
        self.c = c

    def __call__(self, memory: ToTDFSMemory) -> Tuple[str, ...]:
        next_thought = memory.top()
        parent_thought = memory.top_parent()
        validity = (
            ThoughtValidity.VALID_INTERMEDIATE
            if next_thought is None
            else next_thought.validity
        )

        # 1 if the current partial solution is invalid, backtrack to the parent
        # thought.
        if validity == ThoughtValidity.INVALID:
            memory.pop()
            next_thought = memory.top()
            if next_thought and len(next_thought.children) >= self.c:
                memory.pop()

        # 2 if the current partial solution is valid but C children were
        # explored and yet failed to find a final solution, backtrack to the
        # parent thought.
        elif (
            validity == ThoughtValidity.VALID_INTERMEDIATE
            and parent_thought
            and len(parent_thought.children) >= self.c
        ):
            memory.pop(2)

        return tuple(thought.text for thought in memory.current_path())