from typing import Annotated, List, TypedDict, Literal, Dict
from pydantic import BaseModel, Field
import operator


class Section(BaseModel):
    name: str = Field(
        description="章节的标题，应简洁明了地概括本章节的主题或主旨。",
    )
    description: str = Field(
        description="简要说明这一章节将围绕哪些内容进行展开，字数在50-100字。描述应突出章节的写作重点，明确要讨论的问题、角度或结构。",
    )
    research: bool = Field(
        description="判断该章节的写作是否需要进行网络搜索。若该部分需要引用外部资料、数据、案例或是知识储备不足以独立完成，则应标记为需要研究。"
    )
    content: str = Field(
        description="章节的主要写作内容。暂时留空，后续将用于填写具体文本内容"
    )


class Sections(BaseModel):
    sections: List[Section] = Field(
        description="包含报告的各个章节",
    )


class SearchQuery(BaseModel):
    search_query: str = Field(None, description="网络搜索的查询")


class Queries(BaseModel):
    thought: str = Field(None, description="思考过程")
    queries: List[SearchQuery] = Field(
        description="包含网络搜索的查询列表",
    )


'''作为整张图的状态输入输出，输入是报告主题，输出是最终报告'''


class ReportStateInput(TypedDict):
    topic: str  # 报告主题


class ReportStateOutput(TypedDict):
    final_report: str  # 最终报告


'''管理报告的状态，有报告的各个章节'''


class ReportState(TypedDict):
    topic: str
    feedback_on_report_plan: str
    sections: list[Section]
    # 维护一个推理链便于后面模型思考怎么做
    thinking_process: Annotated[list[str], operator.add]
    completed_sections: Annotated[list, operator.add]
    current_section_index: int  # 目前正在研究的章节序号
    final_report: str  # Final report
    search_queries: Annotated[Dict[str, List[SearchQuery]], operator.or_]
    # 键是来源，值是title,url等为键的字典
    sources: Annotated[Dict[str, Dict[str, str]], operator.or_]
    source_str: str