"""State definitions for the Deep Research agent workflow."""

import operator
from typing import Annotated, Optional, List, Dict, Any
from langchain_core.messages import MessageLikeRepresentation
from langgraph.graph import MessagesState
from pydantic import BaseModel, Field
from typing_extensions import TypedDict


###################
# Structured Outputs
###################

class ConductResearch(BaseModel):
    """Call this tool to conduct research on a specific topic."""
    research_topic: str = Field(
        description="The topic to research. Should be a single topic, described in high detail (at least a paragraph).",
    )

class BacktrackResearch(BaseModel):
    """Call this tool to backtrack and refine a previous research topic using prior findings."""
    target_topic: str = Field(
        description="The previous topic to refine. Emphasize the aspect of the topic that you want to revisit.",
    )
    reason: str = Field(
        description="Why backtracking is needed (e.g., gaps, contradictions, low confidence, or newly surfaced angle).",
    )
    previous_related_information: str = Field(
        description="Concise excerpt of prior findings relevant to this backtrack. Provide only the most pertinent passages.",
    )
    updated_focus: str = Field(description="Refined focus or hypotheses to verify during backtracking.")

class ResearchComplete(BaseModel):
    """Call this tool to indicate that the research is complete."""
    
class ClarifyWithUser(BaseModel):
    """Model for user clarification requests."""
    need_clarification: bool = Field(
        description="Whether the user needs to be asked a clarifying question.",
    )
    question: str = Field(
        description="A question to ask the user to clarify the report scope",
    )
    verification: str = Field(
        description="Verify message that we will start research after the user has provided the necessary information.",
    )

class ResearchQuestion(BaseModel):
    """Research question and brief for guiding research."""
    research_brief: str = Field(
        description="A research question that will be used to guide the research.",
    )


###################
# State Definitions
###################

def override_reducer(current_value, new_value):
    """Reducer function that allows overriding values in state."""
    if isinstance(new_value, dict) and new_value.get("type") == "override":
        return new_value.get("value", new_value)
    else:
        return operator.add(current_value, new_value)

class AgentInputState(MessagesState):
    """InputState is only 'messages'."""

class AgentState(MessagesState):
    """Main agent state containing messages and research data."""
    supervisor_messages: Annotated[List[MessageLikeRepresentation], override_reducer]
    research_brief: Optional[str] = None
    raw_notes: Annotated[List[str], override_reducer] = []
    notes: Annotated[List[str], override_reducer] = []
    final_report: Optional[str] = None
    # Persist the exact prompt used to generate the final report
    final_report_prompt: Optional[str] = None
    # Track research topics issued each round by the supervisor
    research_topics_by_round: Annotated[List[dict], override_reducer] = []

class SupervisorState(TypedDict):
    """State for the supervisor that manages research tasks."""
    supervisor_messages: Annotated[List[MessageLikeRepresentation], override_reducer]
    research_brief: str
    notes: Annotated[List[str], override_reducer]
    research_iterations: int
    raw_notes: Annotated[List[str], override_reducer]
    # Accumulates entries like {"round": int, "topics": [str, ...]}
    research_topics_by_round: Annotated[List[dict], override_reducer]

class ResearcherState(TypedDict):
    """State for individual researchers conducting research."""
    researcher_messages: Annotated[List[MessageLikeRepresentation], operator.add]
    tool_call_iterations: int
    research_topic: str
    compressed_research: str
    raw_notes: Annotated[List[str], override_reducer]
    supervisor_messages: Annotated[List[MessageLikeRepresentation], override_reducer]
    # Optional backtracking context when invoked via BacktrackResearch
    backtracking_context: Optional[Dict[str, Any]]

class ResearcherOutputState(BaseModel):
    """Output state from individual researchers."""
    compressed_research: str
    raw_notes: Annotated[List[str], override_reducer] = []
