# state.py
from typing import TypedDict, List, Dict, Optional, Any, Annotated
from langgraph.graph.message import add_messages
from datetime import datetime

class AgentState(TypedDict):
    agent_id: str
    role: str
    last_action: str
    last_updated: str
    status: str  # "idle", "working", "blocked", "done"
    output: Optional[Any]

class SharedContext(TypedDict):
    # 任務(wù)元信息
    session_id: str
    task_id: str
    task_goal: str
    task_constraints: List[str]
    current_phase: str
    phase_history: List[Dict[str, Any]]

    # 共享數(shù)據(jù)區(qū)
    shared_memory: Dict[str, Any]

    # Agent 狀態(tài)池
    agent_states: Dict[str, AgentState]

    # 上下文控制
    context_version: Annotated[int, lambda x, y: x + 1]  # 自動(dòng)遞增版本號(hào)
    event_queue: Annotated[List[Dict], add_messages]      # 事件隊(duì)列（LangGraph 原生支持）

    # 時(shí)間戳
    created_at: str
    updated_at: str

[Researcher Agent]
       │
       ▼
[更新 state.event_queue] → LangGraph 內(nèi)部流轉(zhuǎn)（Writer 下一步消費(fèi)）
       │
       ▼
[發(fā)布到 Redis "agent_events"] → 外部系統(tǒng)實(shí)時(shí)響應(yīng)
       │
       ├→ [前端儀表盤] 實(shí)時(shí)顯示“研究員已獲取數(shù)據(jù)”
       ├→ [日志服務(wù)] 寫入結(jié)構(gòu)化日志
       └→ [告警服務(wù)] 檢測(cè)異常事件（如 confidence < 0.5）

# agents/researcher.py
from event_bus.redis_adapter import RedisEventBus

# 全局 EventBus 實(shí)例（生產(chǎn)環(huán)境建議用依賴注入或 Singleton）
redis_bus = RedisEventBus()

def researcher_agent(state: SharedContext) -> Dict[str, Any]:
    data_snippet = "AI市場(chǎng)規(guī)模2025年預(yù)計(jì)達(dá)$3000億"
    confidence_score = 0.92

    # 寫入 LangGraph 內(nèi)部 event_queue（用于狀態(tài)機(jī)流轉(zhuǎn)）
    internal_event = {
        "event_id": str(uuid.uuid4()),
        "event_type": "context.update",
        "source": "researcher",
        "target": "writer",
        "payload": {
            "key": "market_data",
            "value": data_snippet,
            "confidence": confidence_score
        }
    }

    # 同時(shí)發(fā)布到 Redis（用于外部系統(tǒng)監(jiān)聽）
    redis_bus.publish("agent_events", {
        "event_type": "researcher.data_fetched",
        "source_agent": "researcher",
        "payload": {
            "data_snippet": data_snippet,
            "confidence": confidence_score,
            "task_id": state["task_id"]
        },
        "context_version": state["context_version"]
    })

    return {
        "shared_memory": {**state["shared_memory"], "latest_data": data_snippet},
        "event_queue": [internal_event],  # 供 LangGraph 內(nèi)部消費(fèi)
        "agent_states": { ... }
    }

# monitor/realtime_logger.py
import json
from event_bus.redis_adapter import RedisEventBus

def on_agent_event(message):
    data = json.loads(message['data'])
    event_type = data["event_type"]

    if event_type == "researcher.data_fetched":
        print(f" 實(shí)時(shí)監(jiān)控: 研究員獲取數(shù)據(jù): {data['payload']['data_snippet']}")
    elif event_type == "report.draft_ready":
        print(f" 報(bào)告已生成，長度: {data['payload']['length']} 字符")

# 啟動(dòng)監(jiān)聽
bus = RedisEventBus()
bus.subscribe("agent_events", on_agent_event)

# checkpoint/event_trigger.py
from langgraph.checkpoint.base import BaseCheckpointSaver

class EventEmittingCheckpointer(BaseCheckpointSaver):
    def put(self, config, checkpoint, metadata):
        # 先調(diào)用原生保存邏輯
        super().put(config, checkpoint, metadata)

        # 發(fā)布“狀態(tài)已保存”事件到 EventBus
        redis_bus.publish("system_events", {
            "event_type": "checkpoint.saved",
            "thread_id": config["configurable"]["thread_id"],
            "version": checkpoint["channel_versions"]["__root__"],
            "timestamp": metadata["write_timestamp"]
        })

def get_role_view(state: SharedContext, role: str) -> Dict[str, Any]:
    view_templates = {
        "planner": ["task_goal", "current_phase", "agent_states"],
        "researcher": ["task_goal", "shared_memory.queries", "task_constraints"],
        "writer": ["shared_memory.latest_data", "task_goal"]
    }

    view = {}
    for key in view_templates.get(role, []):
        if '.' in key:
            parts = key.split('.')
            val = state
            for part in parts:
                val = val.get(part, {}) if isinstance(val, dict) else {}
            view[key] = val
        else:
            view[key] = state.get(key)
    return view

# 在 researcher_agent 中調(diào)用
view = get_role_view(state, "researcher")
print(f"[Researcher View] {view}")

# graph_builder.py
from langgraph.graph import StateGraph
from langgraph.checkpoint.sqlite import SqliteSaver
from .state import SharedContext
from .agents import planner_agent, researcher_agent, writer_agent

# 初始化檢查點(diǎn)（持久化 + 版本歷史）
memory = SqliteSaver.from_conn_string("checkpoints.db")

def build_graph():
    workflow = StateGraph(SharedContext)

    workflow.add_node("planner", planner_agent)
    workflow.add_node("researcher", researcher_agent)
    workflow.add_node("writer", writer_agent)

    workflow.set_entry_point("planner")
    workflow.add_edge("planner", "researcher")
    workflow.add_edge("researcher", "writer")

    # 編譯時(shí)綁定檢查點(diǎn)
    return workflow.compile(checkpointer=memory)

# agents/planner.py
def planner_agent(state: SharedContext) -> Dict[str, Any]:
    new_phase = "data_collection"

    # 當(dāng) phase 改變時(shí)，強(qiáng)制創(chuàng)建快照（用于回滾）
    if state["current_phase"] != new_phase:
        # 在返回值中不直接支持“創(chuàng)建快照”，但可通過外部機(jī)制或自定義 reducer 實(shí)現(xiàn)
        # 這里我們通過 event 通知外部系統(tǒng)創(chuàng)建快照
        snapshot_event = {
            "event_id": str(uuid.uuid4()),
            "event_type": "snapshot.request",
            "reason": f"phase_changed: {state['current_phase']} → {new_phase}",
            "manual_trigger": True
        }

        return {
            "current_phase": new_phase,
            "event_queue": [snapshot_event],
            # ... 其他字段
        }

# rollback_manager.py
from .graph_builder import app

def rollback_to_version(thread_id: str, target_version: int):
    """回滾到指定版本"""
    # 獲取所有歷史狀態(tài)
    history = list(app.get_state_history({"configurable": {"thread_id": thread_id}}))

    for snapshot in reversed(history):
        if snapshot.values["context_version"] == target_version:
            # 執(zhí)行回滾
            app.update_state(
                {"configurable": {"thread_id": thread_id}},
                snapshot.values,  # 完整狀態(tài)覆蓋
                as_node="system_rollback"  # 虛擬節(jié)點(diǎn)名
            )
            print(f"已回滾到版本 {target_version}")
            return True
    return False

def detect_context_drift(current_state: SharedContext, baseline: Dict) -> bool:
    """檢測(cè)關(guān)鍵字段是否被外部修改"""
    drift_fields = ["task_goal", "task_constraints"]
    for field in drift_fields:
        if current_state.get(field) != baseline.get(field):
            return True
    return False


if detect_context_drift(latest_state, original_baseline):
    # 觸發(fā) Planner 重新規(guī)劃
    app.update_state(
        {"configurable": {"thread_id": thread_id}},
        {"event_queue": [{
            "event_id": "drift_detected",
            "event_type": "system.reset",
            "payload": {"reason": "task_goal_modified"}
        }]},
        as_node="planner"  # 強(qiáng)制在 planner 上下文中處理
    )

context["user_feedback"] = get_latest_feedback()
context["market_data"] = fetch_real_time_data()