from typing import List, Annotated, Literal
from typing_extensions import TypedDict
import operator
from pydantic import BaseModel, Field
from langchain_core.documents import Document

# --- Pydantic Models for Structured Output ---

class ResearchPlan(BaseModel):
     """Structured research plan."""
     background: str = Field(description="Brief background context for the research topic.")
     key_questions: List[str] = Field(description="List of 5 specific research questions the plan aims to address.")
     search_strategy_summary: str = Field(description="A brief summary of the types of studies or concepts to search for.")

class SearchQueries(BaseModel):
    """Structured list of search queries optimized for literature databases."""
    queries: List[str] = Field(
        description="List of 5 optimized search queries based on the research plan, suitable for PubMed or similar databases (e.g., using keywords)."
    )

# --- Type definition for LLM-generated answers ---
class LLMAnsweredQuestion(TypedDict):
    question: str
    comprehensive_answer: str

# --- LangGraph State Definition ---

class MetaAnalysisState(TypedDict):
    """Represents the state of our meta-analysis assistant graph."""
    research_topic: str           # Input topic from the researcher
    research_plan: ResearchPlan   # Generated plan
    search_queries: List[str]     # Optimized queries for retrieval
    retrieved_docs: Annotated[List[Document], operator.add] # Aggregated documents from all queries
    final_conclusion: str         # Final synthesized conclusion

    # --- New fields for the internal knowledge path ---
    use_internal_knowledge: bool  # Flag to indicate if LLM's internal knowledge should be used
    llm_generated_answers: List[LLMAnsweredQuestion] # Answers generated by LLM for key questions
    
    # Enhanced llm_knowledge route fields
    additional_questions: List[str] | None # Additional sub-questions generated if initial conclusion is inadequate
    additional_answers: List[LLMAnsweredQuestion] | None # Answers to additional sub-questions
    conclusion_evaluation_score: int | None # Score 0-5 for conclusion-topic match assessment
    conclusion_evaluation_feedback: str | None # Detailed feedback on conclusion quality
    is_second_iteration: bool | None # Flag to track if we're in the second iteration to prevent infinite loops
    # --- End new fields ---

    # --- New fields for direct reference text input ---
    target_reference_text: str | None # Optional direct reference text for synthesis
    synthesis_input_source: Literal["retrieved_docs", "llm_knowledge", "target_text", "target_text_suitability"] # Specifies the source for synthesis
    # --- End new fields ---

    # --- New field for suitability assessment ---
    suitability_score: int | None  # Score 0-5 for target text suitability assessment
    suitability_assessment: str | None # Detailed assessment explanation
    original_conclusion: str | None # Original human-written conclusion for target_text_suitability assessment
    # --- End new field ---

    # Temporary state for parallel execution (managed by LangGraph)
    current_query: str            # The query currently being processed by a branch
    current_docs: List[Document]  # Docs retrieved for the current_query
