"""
Step 02: State 状态管理 - 主程序

运行方式：
    cd step_02_state
    python main.py
"""

from concept import (
    JaspersoftAgentState,
    create_initial_state,
    update_state,
    save_state_snapshot,
    restore_state_snapshot,
)


def main():
    """演示状态管理的完整使用流程"""

    print("=" * 70)
    print(" Step 02: 理解 State - 状态管理")
    print("=" * 70)
    print()

    # ═══════════════════════════════════════════════════════════════════════════════
    # 场景：用户生成并修改报表的完整流程
    # ═══════════════════════════════════════════════════════════════════════════════

    print("📦 场景：用户生成并修改报表的完整流程")
    print("-" * 70)
    print("""
假设用户执行以下操作：
    1. 描述需求：生成一个销售报表
    2. Agent 生成报表
    3. Agent 验证通过
    4. 用户要求修改标题
    5. 用户要求撤销修改

我们需要状态来追踪这个完整的流程。
""")

    # ═══════════════════════════════════════════════════════════════════════════════
    # 步骤 1：创建初始状态
    # ═══════════════════════════════════════════════════════════════════════════════

    print("\n📋 步骤 1: 创建初始状态")
    print("-" * 40)

    state = create_initial_state("session_001")
    print(f"  会话 ID: {state['session_id']}")
    print(f"  创建时间: {state['created_at']}")
    print(f"  初始状态: {state['status']}")

    # ═══════════════════════════════════════════════════════════════════════════════
    # 步骤 2：用户描述需求
    # ═══════════════════════════════════════════════════════════════════════════════

    print("\n📝 步骤 2: 用户描述需求")
    print("-" * 40)

    state["user_input"] = "生成一个销售报表，显示月度汇总"
    state["intent"] = "initial_generation"
    state["stage"] = "initial"

    print(f"  用户输入: {state['user_input']}")
    print(f"  意图: {state['intent']}")
    print(f"  阶段: {state['stage']}")

    # ═══════════════════════════════════════════════════════════════════════════════
    # 步骤 3：生成报表（模拟）
    # ═══════════════════════════════════════════════════════════════════════════════

    print("\n🔧 步骤 3: Agent 生成报表")
    print("-" * 40)

    # 模拟生成过程
    state["stage"] = "generation"
    generated_jrxml = '''<?xml version="1.0" encoding="UTF-8"?>
<jasperReport name="SalesReport">
    <title>月度销售汇总报表</title>
    <queryString>SELECT product, SUM(amount) FROM sales GROUP BY product</queryString>
    <field name="product" class="java.lang.String"/>
    <field name="amount" class="java.math.BigDecimal"/>
    <band height="100">
        <staticText><text>产品</text></staticText>
        <textField><textFieldExpression>$F{product}</textFieldExpression></textField>
    </band>
</jasperReport>'''

    state["current_jrxml"] = generated_jrxml
    print(f"  生成完成!")
    print(f"  JRXML 长度: {len(generated_jrxml)} 字符")
    print(f"  当前阶段: {state['stage']}")

    # ═══════════════════════════════════════════════════════════════════════════════
    # 步骤 4：验证通过
    # ═══════════════════════════════════════════════════════════════════════════════

    print("\n✅ 步骤 4: 验证通过")
    print("-" * 40)

    # 保存状态快照（修改前的备份）
    state_snapshot = save_state_snapshot(state)
    print(f"  ✓ 保存状态快照")
    print(f"    快照时间: {state_snapshot['timestamp']}")
    print(f"    包含内容: current_jrxml, conversation_history 等")

    # 验证通过，更新状态
    state["status"] = "success"
    state["final_jrxml"] = state["current_jrxml"]
    state["stage"] = "completed"

    print(f"  验证状态: {state['status']}")
    print(f"  最终版本: ✓ 已保存")

    # ═══════════════════════════════════════════════════════════════════════════════
    # 步骤 5：用户要求修改标题
    # ═══════════════════════════════════════════════════════════════════════════════

    print("\n✏️ 步骤 5: 用户要求修改标题")
    print("-" * 40)

    # 保存修改前的快照
    pre_modify_snapshot = save_state_snapshot(state)
    print(f"  ✓ 修改前保存快照")

    # 执行修改
    new_jrxml = state["current_jrxml"].replace("月度销售汇总报表", "2024年销售汇总报表")
    state["current_jrxml"] = new_jrxml
    state["intent"] = "modify_report"
    state["stage"] = "modification"

    print(f"  修改内容: 标题从'月度销售汇总报表'改为'2024年销售汇总报表'")
    print(f"  当前意图: {state['intent']}")

    # ═══════════════════════════════════════════════════════════════════════════════
    # 步骤 6：用户撤销修改
    # ═══════════════════════════════════════════════════════════════════════════════

    print("\n↩️ 步骤 6: 用户撤销修改")
    print("-" * 40)

    # 恢复到修改前的状态
    state = restore_state_snapshot(state, pre_modify_snapshot)
    print(f"  ✓ 撤销成功!")
    print(f"  恢复标题: {state['current_jrxml'][:50]}...")
    print(f"  当前意图: {state['intent']}")

    # ═══════════════════════════════════════════════════════════════════════════════
    # 展示完整状态
    # ═══════════════════════════════════════════════════════════════════════════════

    print("\n\n" + "=" * 70)
    print("📊 完整状态一览")
    print("=" * 70)

    key_fields = [
        ("session_id", "会话ID"),
        ("status", "状态"),
        ("intent", "意图"),
        ("stage", "阶段"),
        ("user_input", "用户输入"),
        ("current_jrxml", "当前JRXML"),
        ("final_jrxml", "最终JRXML"),
        ("created_at", "创建时间"),
        ("updated_at", "更新时间"),
    ]

    for field, desc in key_fields:
        value = state.get(field, "")
        if field in ["current_jrxml", "final_jrxml"] and value:
            value = f"{value[:50]}..." if len(value) > 50 else value
        print(f"  {desc}: {value}")

    # ═══════════════════════════════════════════════════════════════════════════════
    # 总结
    # ═══════════════════════════════════════════════════════════════════════════════

    print("\n\n" + "=" * 70)
    print(" ✅ Step 02 完成!")
    print("=" * 70)
    print("""
学到的关键概念：
    1. State 是 Agent 的"记忆"，在多步骤任务中保持信息
    2. 使用 TypedDict 定义状态，有类型提示更安全
    3. 状态快照用于"撤销"功能
    4. 不同字段用于不同目的：业务数据、对话历史、元信息

下一步：
    继续 Step 03，学习如何把 Tool + State 组合成简单的 Agent
""")


if __name__ == "__main__":
    main()