"""Define the state structures for the agent."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import List, Optional, Any
from sqlalchemy.ext.asyncio import AsyncSession
from app.services.streaming_service import StreamingService

@dataclass
class State:
    """Defines the dynamic state for the agent during execution.

    This state tracks the database session and the outputs generated by the agent's nodes.
    See: https://langchain-ai.github.io/langgraph/concepts/low_level/#state
    for more information.
    """
    # Runtime context (not part of actual graph state)
    db_session: AsyncSession
    
    # Streaming service
    streaming_service: StreamingService
    
    chat_history: Optional[List[Any]] = field(default_factory=list)
    
    reformulated_query: Optional[str] = field(default=None)
    # Using field to explicitly mark as part of state
    answer_outline: Optional[Any] = field(default=None)
    
    # OUTPUT: Populated by agent nodes
    # Using field to explicitly mark as part of state
    final_written_report: Optional[str] = field(default=None)