from typing import TypedDict
from langgraph.graph import StateGraph, END


class BlogState(TypedDict):
    topic: str          # 初始输入主题
    outline: str        # 第一步产出：大纲
    draft: str          # 第二步产出：初稿
    final_article: str  # 第三步产出：最终定稿

def generate_outline_node(state: BlogState) -> dict:
    """第一步：规划大纲"""
    topic = state["topic"]
    
    # 实际场景中这里调用 llm.invoke(...)
    prompt = f"请为主题【{topic}】规划 3 个核心技术小节大纲。"
    
    outline = (
        f"## 1. 为什么在复杂场景下需要状态机\n"
        f"## 2. LangGraph 的核心抽象与拓扑流\n"
        f"## 3. 生产环境的错误隔离与重试实践"
    )
    return {"outline": outline}


def expand_content_node(state: BlogState) -> dict:
    """第二步：根据大纲扩写正文初稿"""
    topic = state["topic"]
    outline = state["outline"]
    
    draft_content = f"# 深入理解 {topic}\n\n"
    sections = outline.strip().split("\n")
    for sec in sections:
        draft_content += f"{sec}\n这里是针对该小节的详细架构解析、代码片段与工程论证...\n\n"
        
    return {"draft": draft_content.strip()}


def polish_article_node(state: BlogState) -> dict:
    """第三步：润色正文与格式整理"""
    draft = state["draft"]
    
    # 模拟润色处理：统一标点、排版代码块并添加尾注
    final_article = draft + "\n\n---\n*注：本文由自动化工作流生成并完成规范性校验。*"
    return {"final_article": final_article}

workflow = StateGraph(BlogState)

# 注册三个节点
workflow.add_node("outline", generate_outline_node)
workflow.add_node("expand", expand_content_node)
workflow.add_node("polish", polish_article_node)

# 设置线性边
workflow.set_entry_point("outline")
workflow.add_edge("outline", "expand")
workflow.add_edge("expand", "polish")
workflow.add_edge("polish", END)

app = workflow.compile()

initial_input = {
    "topic": "LangGraph 状态机编排",
    "outline": "",
    "draft": "",
    "final_article": "",
}

result = app.invoke(initial_input)
print("=== 最终产物 ===")
print(result["final_article"])

for step in app.stream(initial_input):
    for node_name, state_update in step.items():
        print(f"\n>>>> 节点 [{node_name}] 执行完毕，产生状态变更：")
        for k, v in state_update.items():
            preview = v.replace("\n", " ")[:60]
            print(f"  - {k}: {preview}...")

from typing import TypedDict, Literal
from langgraph.graph import StateGraph, END


class RobustBlogState(TypedDict):
    topic: str
    outline: str
    draft: str
    final_article: str
    outline_retries: int  # 重试计数器


def generate_outline(state: RobustBlogState) -> dict:
    retries = state.get("outline_retries", 0)
    topic = state["topic"]
    
    # 模拟第一次生成可能不合格，重试后合格
    if retries == 0:
        bad_outline = "## 1. 简单介绍"
        return {"outline": bad_outline, "outline_retries": retries + 1}
    else:
        good_outline = (
            "## 1. 背景与核心问题\n"
            "## 2. 核心架构设计\n"
            "## 3. 生产部署与调优"
        )
        return {"outline": good_outline, "outline_retries": retries + 1}


def validate_outline_gate(state: RobustBlogState) -> Literal["expand", "retry", "abort"]:
    """路由器：判断大纲章节数与重试上限"""
    outline = state.get("outline", "")
    section_count = outline.count("##")
    retries = state.get("outline_retries", 0)
    
    if section_count >= 3:
        return "expand"
    
    if retries <= 2:
        print(f"[警告] 大纲小节数不足（仅 {section_count} 个），触发第 {retries} 次重试...")
        return "retry"
    
    print("[错误] 大纲重试超限，任务终止")
    return "abort"


# 组装图
workflow = StateGraph(RobustBlogState)
workflow.add_node("outline", generate_outline)
workflow.add_node("expand", expand_content_node)
workflow.add_node("polish", polish_article_node)

workflow.set_entry_point("outline")

# 添加条件边实现自旋纠错
workflow.add_conditional_edges(
    "outline",
    validate_outline_gate,
    {
        "expand": "expand",
        "retry": "outline",  # 指向自身，形成重试环路
        "abort": END,
    }
)

workflow.add_edge("expand", "polish")
workflow.add_edge("polish", END)

app = workflow.compile()

# 伪代码：在节点内部平滑切换模型
def generate_outline_with_llm(state: BlogState) -> dict:
    topic = state["topic"]
    prompt = f"请为【{topic}】构思一篇技术架构深度博客的大纲，包含3-4个核心章节。"
    
    # 这里可以使用任意兼容 LangChain 接口的 LLM
    # response = model.invoke([HumanMessage(content=prompt)])
    # return {"outline": response.content}