"""
schemas.py — Data models and LangGraph State for Academic Deep Researcher.

Contains:
  - Pydantic models for structured LLM outputs
  - TypedDict State definition with reducers
"""

from typing import Annotated, TypedDict, List, Optional, Any
from langgraph.graph.message import add_messages
from pydantic import BaseModel, Field
import operator


# ---------------------------------------------------------------------------
# Structured output schemas (used with LLM.with_structured_output())
# ---------------------------------------------------------------------------

class ClarifyingQA(BaseModel):
    """A single clarifying question and the user's answer."""
    question: str
    answer: str


class PlannerOutput(BaseModel):
    """Structured output from the task_planner node."""
    refined_research_question: str = Field(
        description="One clear standalone research question merging the topic and user clarifications"
    )
    research_tasks: List[str] = Field(
        description="List of specific research sub-tasks")
    semantic_scholar_query: str = Field(
        description="Formal academic query optimised for Semantic Scholar")
    arxiv_query: str = Field(
        description="Technical query optimised for ArXiv preprint search")
    tavily_query: str = Field(
        description="Natural language query for recent web content via Tavily")
    pubmed_query: Optional[str] = Field(
        default=None,
        description="MeSH-term query for PubMed — only set if topic is biomedical"
    )
    needs_pubmed: bool = Field(
        default=False,
        description="True only if the topic is medical, biological, clinical, or pharmacological"
    )


class ResearchPoint(BaseModel):
    """A single rated research finding."""
    claim: str = Field(description="The key finding or fact in 1–2 sentences")
    source_url: str = Field(description="URL of the source")
    source_title: str = Field(description="Title of the source document")
    source_type: str = Field(
        description="One of: paper | preprint | web | news")
    citation_count: Optional[int] = Field(
        default=None, description="Citation count if available from Semantic Scholar")
    confidence_score: float = Field(
        description="0.0–1.0 reliability: peer-reviewed ≤1.0, preprint ≤0.85, web ≤0.7")
    relevance_score: float = Field(
        description="0.0–1.0 relevance to the research topic")


class RaterOutput(BaseModel):
    """Structured output from the rater node."""
    rated_points: List[ResearchPoint] = Field(
        description="Only include points where confidence ≥ 0.4 AND relevance ≥ 0.4"
    )


class FinalReport(BaseModel):
    """Structured output from the report_writer node."""
    executive_summary: str = Field(
        description="3-paragraph overview for the stated audience")
    key_facts: List[str] = Field(
        description="8–12 highest-confidence facts as plain strings")
    detailed_research: str = Field(
        description="600–900 word analysis with inline citations as [Source Title]")
    sources: List[str] = Field(
        description="Formatted references: Title — URL (Year)")
    further_reading: List[str] = Field(
        description="5–7 recommendations: next searches, related topics, key authors")


# ---------------------------------------------------------------------------
# LangGraph State
# ---------------------------------------------------------------------------

class State(TypedDict):
    # Conversation history — add_messages reducer appends and deduplicates
    messages: Annotated[List[Any], add_messages]

    # Research setup
    topic: str
    clarifying_questions: List[str]      # generated by clarifier node
    clarifying_qa: List[dict]            # user-approved Q&A pairs
    clarifications_approved: bool

    # Task planning
    refined_topic: str
    research_tasks: List[str]
    # serialised PlannerOutput (.model_dump())
    agent_queries: Optional[dict]

    # Parallel research results
    # operator.add safely accumulates results from all parallel agents
    raw_results: Annotated[List[dict], operator.add]

    # Post-processing
    synthesized_findings: str
    rated_points: List[dict]             # serialised ResearchPoint list

    # Final output
    final_report: Optional[dict]         # serialised FinalReport
    email_sent: bool
    recipient_email: str
