# 每次任务使用唯一的 thread_id
config = {"configurable": {"thread_id": "task-20241215-report-analysis"}}

# 第一次运行：Agent 开始执行
graph.invoke({"task": "分析Q3财务报告"}, config=config)

# 服务重启后，使用相同的 thread_id 恢复
graph.invoke(None, config=config)  # 从上次停止的地方继续

from typing import TypedDict, Annotated
import operator
import time
from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, SystemMessage
from langgraph.graph import StateGraph, START, END
from langgraph.checkpoint.memory import MemorySaver
# 生产环境替换为：
# from langgraph.checkpoint.postgres import PostgresSaver

# ---- 定义 Agent 状态 ----

class DocumentAnalysisState(TypedDict):
    """
    长文档分析 Agent 的状态。
    每个字段在不同阶段被填充，全部完成时任务结束。
    """
    document_text: str                              # 原始文档内容
    document_chunks: list[str]                      # 切分后的文档片段
    chunk_analyses: Annotated[list, operator.add]   # 每个片段的分析结果（累加）
    summary: str                                    # 整合摘要
    final_report: str                               # 最终报告
    current_phase: str                              # 当前所在阶段（用于进度展示）
    progress: str                                   # 进度说明（面向用户）


llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

# ---- 阶段1：文档切分 ----

def chunk_document(state: DocumentAnalysisState) -> DocumentAnalysisState:
    """
    把长文档切分成小片段。
    为什么要切分？因为单个 LLM 调用的上下文窗口有限，
    一次只能处理一部分文档。
    """
    print("[阶段1] 开始切分文档...")

    text = state["document_text"]
    # 简单的切分策略：每 1000 字为一段（实际生产中应按段落切分）
    chunk_size = 1000
    chunks = [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]

    print(f"[阶段1] 文档切分完成，共 {len(chunks)} 个片段")
    return {
        "document_chunks": chunks,
        "chunk_analyses": [],  # 初始化为空列表
        "current_phase": "chunked",
        "progress": f"文档切分完成，共 {len(chunks)} 个片段，开始逐段分析",
    }


# ---- 阶段2：逐段分析 ----

def analyze_chunks(state: DocumentAnalysisState) -> DocumentAnalysisState:
    """
    逐段分析文档内容。
    这是最耗时的阶段——如果有 50 个片段，需要调用 LLM 50 次。
    每次调用完成后，Checkpointer 会保存当前状态。
    如果这个阶段中途被中断，下次从已分析到的片段继续，不需要重新分析已完成的片段。
    """
    chunks = state["document_chunks"]
    already_analyzed = len(state.get("chunk_analyses", []))

    print(f"[阶段2] 开始逐段分析，共 {len(chunks)} 段，已完成 {already_analyzed} 段")

    # 注意：这里只分析尚未处理的片段
    # 如果之前中断过，chunk_analyses 里已有部分结果，直接从断点继续
    new_analyses = []
    for i, chunk in enumerate(chunks[already_analyzed:], start=already_analyzed):
        print(f"  分析第 {i+1}/{len(chunks)} 段...")

        response = llm.invoke([
            SystemMessage(content=(
                "你是一名文档分析专家。请提取以下文档片段的关键信息：\n"
                "1. 主要观点（1-2句）\n"
                "2. 重要数据或事实（如有）\n"
                "3. 潜在风险或问题（如有）\n"
                "输出格式：JSON，字段为 main_point、key_facts、risks"
            )),
            HumanMessage(content=f"文档片段 {i+1}：\n{chunk}"),
        ])

        new_analyses.append({
            "chunk_index": i,
            "analysis": response.content,
        })

        # 模拟耗时操作（实际场景中这里就是 LLM 调用的延迟）
        # 每分析完一段，LangGraph 会在下个节点执行前保存状态
        # 如果此时服务重启，chunk_analyses 里的已有内容会被保留

    total_analyses = already_analyzed + len(new_analyses)
    print(f"[阶段2] 逐段分析完成，共分析 {total_analyses} 段")
    return {
        "chunk_analyses": new_analyses,  # Annotated[list, operator.add] 会自动追加
        "current_phase": "analyzed",
        "progress": f"逐段分析完成，共分析 {total_analyses} 个片段",
    }


# ---- 阶段3：整合摘要 ----

def synthesize_summary(state: DocumentAnalysisState) -> DocumentAnalysisState:
    """
    把所有片段的分析结果整合成一份连贯的摘要。
    """
    print("[阶段3] 开始整合摘要...")

    # 把所有片段分析结果拼接成输入
    all_analyses = "\n\n".join([
        f"片段{item['chunk_index']+1}分析：{item['analysis']}"
        for item in state["chunk_analyses"]
    ])

    response = llm.invoke([
        SystemMessage(content=(
            "基于以下各片段的分析结果，生成一份整合摘要。\n"
            "要求：综合所有主要观点、汇总重要数据、列出全部风险点。\n"
            "字数：300-500字。"
        )),
        HumanMessage(content=all_analyses),
    ])

    print("[阶段3] 摘要整合完成")
    return {
        "summary": response.content,
        "current_phase": "summarized",
        "progress": "摘要整合完成，准备生成最终报告",
    }


# ---- 阶段4：生成最终报告 ----

def generate_report(state: DocumentAnalysisState) -> DocumentAnalysisState:
    """
    基于摘要生成结构化的最终报告。
    """
    print("[阶段4] 开始生成最终报告...")

    response = llm.invoke([
        SystemMessage(content=(
            "基于文档摘要，生成一份专业的分析报告。\n"
            "报告结构：执行摘要、主要发现、风险评估、建议措施。\n"
            "格式：Markdown，使用标题和列表。"
        )),
        HumanMessage(content=f"文档摘要：\n{state['summary']}"),
    ])

    print("[阶段4] 最终报告生成完成")
    return {
        "final_report": response.content,
        "current_phase": "completed",
        "progress": "任务完成",
    }


# ---- 构建支持断点续跑的图 ----

def build_resumable_graph(checkpointer):
    """
    构建带 Checkpointer 的分析图。
    checkpointer 参数决定状态存储在哪里：
      - MemorySaver()：存内存（开发用）
      - SqliteSaver.from_conn_string("./agent_state.db")：存本地文件
      - PostgresSaver(...)：存数据库（生产用）
    """
    builder = StateGraph(DocumentAnalysisState)

    builder.add_node("chunk_document", chunk_document)
    builder.add_node("analyze_chunks", analyze_chunks)
    builder.add_node("synthesize_summary", synthesize_summary)
    builder.add_node("generate_report", generate_report)

    # 串行管道：四个阶段依次执行
    builder.add_edge(START, "chunk_document")
    builder.add_edge("chunk_document", "analyze_chunks")
    builder.add_edge("analyze_chunks", "synthesize_summary")
    builder.add_edge("synthesize_summary", "generate_report")
    builder.add_edge("generate_report", END)

    # 关键：传入 checkpointer，启用状态持久化
    return builder.compile(checkpointer=checkpointer)


# ---- 演示：启动任务、模拟中断、断点续跑 ----

def demo_resumable_task():
    """
    完整演示：
    1. 启动一个文档分析任务
    2. 用相同 thread_id 查询进度（模拟中断后恢复）
    """
    checkpointer = MemorySaver()  # 生产环境换成 PostgresSaver
    graph = build_resumable_graph(checkpointer)

    # 每次任务使用固定且唯一的 thread_id
    # 实际场景中可以是 "用户ID-任务ID" 的组合
    thread_id = "user-001-q3-report-analysis"
    config = {"configurable": {"thread_id": thread_id}}

    sample_document = """
    第一章：市场概述
    本季度市场总规模达到 500 亿元，同比增长 23%。
    主要增长驱动力来自 AI 相关产品和服务，占新增市场的 67%。
    ...（模拟长文档内容）...
    第二章：竞争格局
    头部三家企业占据市场份额的 78%，格局相对集中。
    新进入者增加，竞争加剧，价格战风险上升。
    """

    print("=== 启动文档分析任务 ===")
    final_state = graph.invoke(
        {
            "document_text": sample_document,
            "document_chunks": [],
            "chunk_analyses": [],
            "summary": "",
            "final_report": "",
            "current_phase": "started",
            "progress": "任务已创建",
        },
        config=config,
    )
    print(f"\n任务完成！最终阶段：{final_state['current_phase']}")

    # 用相同的 thread_id 查询历史状态
    history = list(graph.get_state_history(config))
    print(f"\n共保存了 {len(history)} 个检查点")
    for snapshot in history[:3]:  # 只展示最新3个
        phase = snapshot.values.get("current_phase", "unknown")
        print(f"  - 阶段：{phase}，已分析片段数：{len(snapshot.values.get('chunk_analyses', []))}")

    return final_state["final_report"]


demo_resumable_task()

# 生产环境配置：使用 PostgreSQL 存储状态
# 需要安装：pip install langgraph-checkpoint-postgres

from langgraph.checkpoint.postgres import PostgresSaver
import psycopg

# 连接 PostgreSQL（实际使用时替换为真实连接信息）
DB_URL = "postgresql://user:password@localhost:5432/agent_db"

def get_production_checkpointer():
    """
    生产环境 Checkpointer 工厂函数。
    PostgresSaver 会自动创建所需的数据库表。
    """
    conn = psycopg.connect(DB_URL, autocommit=True)
    checkpointer = PostgresSaver(conn)
    checkpointer.setup()  # 创建所需的数据库表（首次使用时执行）
    return checkpointer


# 用户今天的会话
def user_session_day1(thread_id: str):
    checkpointer = get_production_checkpointer()
    graph = build_resumable_graph(checkpointer)
    config = {"configurable": {"thread_id": thread_id}}

    print("用户第一天：开始分析任务...")
    # 假设这里任务执行了一半，用户关闭了浏览器
    # 状态已自动保存到 PostgreSQL


# 用户明天的会话（服务早已重启过多次）
def user_session_day2(thread_id: str):
    checkpointer = get_production_checkpointer()
    graph = build_resumable_graph(checkpointer)
    config = {"configurable": {"thread_id": thread_id}}

    # 查询当前任务状态
    current_state = graph.get_state(config)
    if current_state.values:
        phase = current_state.values.get("current_phase")
        progress = current_state.values.get("progress")
        print(f"找到未完成的任务，当前阶段：{phase}")
        print(f"进度：{progress}")
        print("从断点继续执行...")

        # 传入 None 作为输入，表示"继续已有任务"而非"开始新任务"
        result = graph.invoke(None, config=config)
        return result
    else:
        print("没有找到未完成的任务")
        return None

import datetime

def cleanup_old_checkpoints(conn, retention_days: int = 30):
    """
    清理超过 retention_days 天的已完成任务的中间检查点。
    保留策略：
    - 未完成的任务：保留全部（可能还需要恢复）
    - 已完成的任务：只保留最终状态快照，删除中间步骤
    - 超过 retention_days 的已完成任务：全部删除
    """
    cutoff_date = datetime.datetime.now() - datetime.timedelta(days=retention_days)

    with conn.cursor() as cur:
        # 删除超过保留期的已完成任务的检查点
        # 注意：这里假设"completed"阶段代表任务完成
        # 实际实现需要根据具体的表结构调整 SQL
        cur.execute("""
            DELETE FROM checkpoints
            WHERE created_at < %s
            AND thread_id IN (
                SELECT DISTINCT thread_id FROM checkpoints
                WHERE checkpoint_data->>'current_phase' = 'completed'
            )
        """, (cutoff_date,))

        deleted_count = cur.rowcount
        conn.commit()

    print(f"清理完成：删除了 {deleted_count} 条旧检查点记录")
    return deleted_count


def keep_only_final_checkpoint(conn, thread_id: str):
    """
    任务完成后，删除该任务的中间检查点，只保留最终状态。
    节省存储空间，同时保留完成记录。
    """
    with conn.cursor() as cur:
        # 找到最新（最终）的检查点 ID
        cur.execute("""
            SELECT checkpoint_id FROM checkpoints
            WHERE thread_id = %s
            ORDER BY created_at DESC
            LIMIT 1
        """, (thread_id,))
        latest = cur.fetchone()

        if latest:
            # 删除除最新检查点外的所有中间检查点
            cur.execute("""
                DELETE FROM checkpoints
                WHERE thread_id = %s AND checkpoint_id != %s
            """, (thread_id, latest[0]))
            conn.commit()
            print(f"任务 {thread_id}：中间检查点已清理，保留最终状态")

import asyncio

async def run_with_progress(graph, state, config, progress_callback):
    """
    运行 Agent 并在每个节点完成时回调进度。
    progress_callback：一个函数，接收进度信息并展示给用户（如推送给前端）
    """
    async for event in graph.astream_events(state, config=config, version="v2"):
        event_type = event["event"]
        node_name = event.get("name", "")

        if event_type == "on_chain_start" and node_name in [
            "chunk_document", "analyze_chunks", "synthesize_summary", "generate_report"
        ]:
            # 节点开始时推送进度
            phase_names = {
                "chunk_document": "正在切分文档...",
                "analyze_chunks": "正在逐段分析（耗时较长，请耐心等待）...",
                "synthesize_summary": "正在整合摘要...",
                "generate_report": "正在生成最终报告...",
            }
            await progress_callback({
                "status": "running",
                "phase": node_name,
                "message": phase_names.get(node_name, f"正在执行：{node_name}"),
            })

        elif event_type == "on_chain_end" and node_name == "generate_report":
            # 最终节点完成时推送完成通知
            await progress_callback({
                "status": "completed",
                "phase": "done",
                "message": "分析报告生成完毕",
            })


# 配合 FastAPI + WebSocket 使用示例（伪代码）
# @app.websocket("/ws/task/{thread_id}")
# async def task_websocket(websocket: WebSocket, thread_id: str):
#     await websocket.accept()
#
#     async def push_progress(progress_data):
#         await websocket.send_json(progress_data)
#
#     config = {"configurable": {"thread_id": thread_id}}
#     await run_with_progress(graph, None, config, push_progress)