import operator

from langgraph.graph import MessagesState
from typing_extensions import Annotated, List


class State(MessagesState):
    summary: str
    question: str
    chunked_last_tool_call: bool
    file_reference: str  # Attachment file reference: a path, URL, or unique ID


class MathAgentState(MessagesState):
    """
    State for the math agent containing message history and research metadata.

    This state tracks the agent's conversation, iteration count for limiting
    tool calls, the research topic being investigated, compressed findings,
    and raw research notes for detailed analysis.
    """
    tool_call_iterations: int
    question: str
    compressed_research: str
    raw_notes: Annotated[List[str], operator.add]


class MathAgentOutputState(MessagesState):
    """
    Output state for the math agent containing final results.

    This represents the final output of the solving process with steps
    and raw notes from the solving process.
    """
    compressed_research: str
    raw_notes: Annotated[List[str], operator.add]