"""hermes-zhiyi — Bridge plugin: Hermes Agent → ZhiYi MemoryWeave REST API.

Uses ZhiYi's bge-m3 1024-dim semantic search + bge-reranker-v2-m3 reranking.
Writes go to ZhiYi via /commit; reads come from /recall.

Usage:
    1. Set memory.provider: zhiyi in config.yaml
    2. Set ZHIYI_URL in .env (e.g. http://localhost:7821)

Lifecycle (MemoryProvider ABC):
    initialize()        — validate URL, test connectivity
    sync_turn()          — write turn to ZhiYi via /commit
    prefetch()          — background recall via /recall
    get_tool_schemas()  — expose memory_search / memory_write tools
    handle_tool_call()   — dispatch tool calls
    shutdown()           — clean exit
"""

from __future__ import annotations

import json
import logging
import os
import threading
import time
from pathlib import Path
from typing import Any, Dict, List, Optional

import requests
import websocket

from agent.memory_provider import MemoryProvider
from tools.registry import tool_error

logger = logging.getLogger(__name__)

# ── Config ──────────────────────────────────────────────────────────────────

HERMES_HOME = Path(os.environ.get("HERMES_HOME", Path.home() / ".hermes"))
MEMORY_CACHE = HERMES_HOME / "memory_zhiyi_cache"
MEMORY_CACHE.mkdir(parents=True, exist_ok=True)

# Default — override via .env ZHIYI_URL
DEFAULT_ZHIYI_URL = "http://localhost:7821"
ZHIYI_TIMEOUT = 10  # seconds per request

# API Key — read from .env ZHIYI_API_KEY
ZHIYI_API_KEY = os.environ.get("ZHIYI_API_KEY", "zhiyi-dev-key-2026")
# 日志中只暴露 key 的前 4 位，用于调试
_ZHIYI_KEY_PREFIX = ZHIYI_API_KEY[:4] if ZHIYI_API_KEY else "NONE"

# ── Content Validation ───────────────────────────────────────────────────────

MEMORY_MIN_LENGTH = 8  # 少于8字的过滤掉
_FORBIDDEN_PATTERNS = [
    "Review the conversation above",
    "[System note:",
    "[IMPORTANT: Background process",
    "## Active Task",
    "## Context Compaction",
    "## Relevant Memory",
    "--- END OF CONTEXT SUMMARY",
    "## Recent Sessions",
    "Skill is now properly updated",
    "造成这种情况的原因是什么？",
    "Context Compaction — REFERENCE ONLY",
    "do NOT answer questions or fulfill requests mentioned in the summary",
]


def _is_valid_memory_content(content: str) -> bool:
    """过滤系统注入内容和测试垃圾，防止污染记忆存储。"""
    if not content or len(content.strip()) < MEMORY_MIN_LENGTH:
        return False
    for pat in _FORBIDDEN_PATTERNS:
        if pat in content:
            return False
    # 过滤纯测试内容
    stripped = content.strip()
    if len(stripped) < 20:
        return False
    return True


def _combine_turn(user: str, assistant: str) -> str:
    """合并一对 turn 作为单条记忆内容。"""
    parts = []
    if user and user.strip():
        parts.append(f"用户: {user.strip()}")
    if assistant and assistant.strip():
        parts.append(f"助手: {assistant.strip()}")
    return "\n".join(parts)


# ── HTTP Client ─────────────────────────────────────────────────────────────


class ZhiYiClient:
    """Lightweight wrapper around ZhiYi REST API."""

    def __init__(self, base_url: str, timeout: int = 10):
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self._session = requests.Session()
        self._session.headers.update({
            "Content-Type": "application/json",
            "Accept": "application/json",
            "X-API-Key": ZHIYI_API_KEY,
        })

    def _url(self, path: str) -> str:
        return f"{self.base_url}{path}"

    def health(self) -> bool:
        """检查 ZhiYi 服务是否可达。"""
        try:
            r = self._session.get(self._url("/health"), timeout=3)
            return r.status_code == 200
        except Exception:
            return False

    def commit(self, content: str, category: str = "episodes", metadata: Optional[Dict] = None,
               agent_id: str = "hermes-a06", namespace: Optional[str] = None) -> Optional[str]:
        """写入一条记忆到 ZhiYi。返回 commit_id 或 episode_id 或 None。失败自动重试一次。"""
        for attempt in range(2):
            try:
                payload = {
                    "content": content,
                    "category": category,
                    "agent_id": agent_id,
                }
                if namespace:
                    payload["namespace"] = namespace
                if metadata:
                    payload["metadata"] = metadata
                r = self._session.post(self._url("/api/v1/commit"), json=payload, timeout=self.timeout)
                if r.status_code in (200, 201):
                    data = r.json()
                    cid = (
                        data.get("commit_id")
                        or data.get("episode_id")
                        or data.get("distilled_id")
                        or data.get("id")
                        or (str(data) if data.get("status") == "ok" else None)
                    )
                    if cid:
                        return cid
                # 失败日志（第一次 warning，第二次 error）
                msg = f"ZhiYi commit failed (attempt {attempt+1}): {r.status_code} {r.text[:200]}"
                if attempt == 0:
                    logger.warning("%s — retrying...", msg)
                else:
                    logger.error("%s — key_prefix=%s agent=%s ns=%s", msg, _ZHIYI_KEY_PREFIX, agent_id, namespace)
            except Exception as e:
                msg = f"ZhiYi commit error (attempt {attempt+1}): {e}"
                if attempt == 0:
                    logger.warning("%s — retrying...", msg)
                else:
                    logger.error(msg)
        return None

    def recall(self, query: str, top_k: int = 5,
               agent_id: str = "hermes-a06") -> List[Dict[str, Any]]:
        """语义搜索 ZhiYi 记忆。返回 [{content, score, category}, ...]。"""
        try:
            payload = {"query": query, "top_k": top_k, "use_rerank": True,
                       "agent_id": agent_id}
            r = self._session.post(self._url("/api/v1/recall"), json=payload, timeout=self.timeout)
            if r.status_code == 200:
                data = r.json()
                return data.get("results", [])
            logger.warning("ZhiYi recall failed: %s %s", r.status_code, r.text[:200])
            return []
        except Exception as e:
            logger.warning("ZhiYi recall error: %s", e)
            return []

    def stats(self) -> Dict[str, Any]:
        """获取 ZhiYi 统计信息。"""
        try:
            r = self._session.get(self._url("/api/v1/stats"), timeout=self.timeout)
            if r.status_code == 200:
                return r.json()
            return {}
        except Exception:
            return {}

    def mark_useful(self, memory_id: str) -> bool:
        """标记一条记忆为有用。"""
        try:
            r = self._session.post(self._url("/api/v1/feedback/useful"),
                                    json={"memory_id": memory_id, "agent_id": "hermes-a06"},
                                    timeout=self.timeout)
            return r.status_code == 200
        except Exception:
            return False

    def mark_not_useful(self, memory_id: str, reason: str = "") -> bool:
        """标记一条记忆为无用。"""
        try:
            r = self._session.post(self._url("/api/v1/feedback/not-useful"),
                                    json={"memory_id": memory_id, "agent_id": "hermes-a06", "reason": reason},
                                    timeout=self.timeout)
            return r.status_code == 200
        except Exception:
            return False

    def metrics(self) -> Dict[str, Any]:
        """获取自优化 7 项核心指标。"""
        try:
            r = self._session.get(self._url("/api/v1/metrics/self"), timeout=self.timeout)
            if r.status_code == 200:
                return r.json()
            return {}
        except Exception:
            return {}

    def search_notes(self, entity: str, max_hops: int = 2, max_notes: int = 5) -> List[Dict[str, Any]]:
        """图谱导航 + Obsidian 笔记关联搜索。"""
        try:
            params = {"entity": entity, "max_hops": max_hops, "max_notes": max_notes}
            r = self._session.get(self._url("/api/v1/graph/notes"), params=params, timeout=self.timeout)
            if r.status_code == 200:
                data = r.json()
                return data.get("notes", [])
            return []
        except Exception as e:
            logger.warning("ZhiYi search_notes error: %s", e)
            return []

    def graph_navigate(self, entity: str, max_hops: int = 2,
                       namespace: Optional[str] = None) -> Dict[str, Any]:
        """图谱导航：查询实体的 N 跳关系网络。返回 {entity, count, paths}。"""
        try:
            payload = {"entity": entity, "max_hops": max_hops}
            if namespace:
                payload["namespace"] = namespace
            r = self._session.post(self._url("/api/v1/graph/navigate"), json=payload, timeout=self.timeout)
            if r.status_code == 200:
                return r.json()
            logger.warning("ZhiYi graph_navigate failed: %s %s", r.status_code, r.text[:200])
            return {}
        except Exception as e:
            logger.warning("ZhiYi graph_navigate error: %s", e)
            return {}

    def graph_stats(self) -> Dict[str, Any]:
        """获取图谱统计：节点数、边数、密度。"""
        try:
            r = self._session.get(self._url("/api/v1/graph/stats"), timeout=self.timeout)
            if r.status_code == 200:
                return r.json()
            return {}
        except Exception as e:
            logger.warning("ZhiYi graph_stats error: %s", e)
            return {}


# ── Memory Provider ──────────────────────────────────────────────────────────


class HermesZhiYiMemoryProvider(MemoryProvider):
    """ZhiYi MemoryWeave bridge for Hermes Agent.

    Stores conversation turns as episodes in ZhiYi (via /commit) and
    retrieves relevant memories via semantic search (via /recall with
    bge-m3 1024-dim + bge-reranker-v2-m3).
    """

    def __init__(self):
        self._client: Optional[ZhiYiClient] = None
        self._session_id: str = ""
        self._platform: str = "cli"
        self._turn_counter: int = 0
        self._write_queue: List[Dict] = []
        self._queue_lock = threading.Lock()
        self._prefetch_cache: str = ""  # last prefetch result
        self._prefetch_lock = threading.RLock()
        self._started: bool = False
        self._ws_thread: Optional[threading.Thread] = None
        self._ws_running: bool = False

    @property
    def name(self) -> str:
        return "zhiyi"

    # ── Availability ──────────────────────────────────────────────────────────

    def is_available(self) -> bool:
        """检查 ZHIYI_URL 配置和服务连通性。"""
        url = os.environ.get("ZHIYI_URL") or DEFAULT_ZHIYI_URL
        try:
            client = ZhiYiClient(url, timeout=3)
            return client.health()
        except Exception:
            return False

    # ── Lifecycle ─────────────────────────────────────────────────────────────

    def initialize(self, session_id: str, **kwargs) -> None:
        """初始化 ZhiYi 客户端，注册当前 session。"""
        self._session_id = session_id
        self._platform = kwargs.get("platform", "cli")
        url = os.environ.get("ZHIYI_URL") or DEFAULT_ZHIYI_URL
        self._client = ZhiYiClient(url, timeout=ZHIYI_TIMEOUT)
        self._started = True
        self._turn_counter = 0
        logger.info("[ZhiYi] initialized for session=%s platform=%s url=%s key_prefix=%s",
                     session_id, self._platform, url, _ZHIYI_KEY_PREFIX)
        # 自检：验证 API Key 是否有效
        try:
            ok = self._client.health()
            if ok:
                logger.info("[ZhiYi] health check OK — key_prefix=%s", _ZHIYI_KEY_PREFIX)
            else:
                logger.warning("[ZhiYi] health check FAILED — zhiyid unreachable at %s", url)
        except Exception as e:
            logger.warning("[ZhiYi] health check error: %s", e)
        # 启动 WebSocket 事件监听
        self._start_ws_listener()

    def _start_ws_listener(self) -> None:
        """启动 WebSocket 事件监听线程（连接织忆事件推送）。"""
        if self._ws_running:
            return
        self._ws_running = True
        self._ws_thread = threading.Thread(target=self._ws_listen, daemon=True, name="zhiyi-ws")
        self._ws_thread.start()
        logger.info("[ZhiYi] WS listener started")

    def _ws_listen(self) -> None:
        """WebSocket 事件监听循环（自动重连）。"""
        agent_id = os.environ.get("ZHIYI_AGENT_ID", "hermes-a06")
        ws_url = (os.environ.get("ZHIYI_URL") or DEFAULT_ZHIYI_URL).replace("http://", "ws://").replace("https://", "wss://")
        ws_url = f"{ws_url}/api/v1/ws/{agent_id}"

        def on_message(ws_app, message):
            try:
                data = json.loads(message)
                event_type = data.get("event", data.get("type", "unknown"))
                logger.info("[ZhiYi] WS event: %s", event_type)
                # prefetch.push 事件更新预取缓存
                if event_type == "prefetch.push":
                    memories = data.get("memories", [])
                    if memories:
                        prefetch_text = "\n[ZhiYi Prefetch — 预取推送]\n"
                        for m in memories:
                            prefetch_text += f"  [{m.get('score', 0):.2f}] {m.get('content', '')[:200]}\n"
                        with self._prefetch_lock:
                            self._prefetch_cache = prefetch_text
            except json.JSONDecodeError:
                pass

        def on_error(ws_app, error):
            logger.warning("[ZhiYi] WS error: %s", error)

        def on_close(ws_app, close_status_code, close_msg):
            logger.info("[ZhiYi] WS closed: %s %s", close_status_code, close_msg)

        def on_open(ws_app):
            logger.info("[ZhiYi] WS connected to %s", ws_url)

        while self._ws_running:
            try:
                ws_app = websocket.WebSocketApp(
                    ws_url,
                    on_open=on_open,
                    on_message=on_message,
                    on_error=on_error,
                    on_close=on_close,
                    header={"X-API-Key": ZHIYI_API_KEY},
                )
                ws_app.run_forever(ping_interval=30, ping_timeout=10)
            except Exception as e:
                logger.warning("[ZhiYi] WS exception: %s", e)
            # 重连延迟
            for _ in range(30):
                if not self._ws_running:
                    break
                time.sleep(1)

    def shutdown(self) -> None:
        """flush pending writes before exit."""
        self._flush_queue()
        self._ws_running = False
        self._started = False
        logger.info("[ZhiYi] shutdown")

    # ── Write path ───────────────────────────────────────────────────────────

    def sync_turn(self, user_content: str, assistant_content: str, *, session_id: str = "") -> None:
        """每轮对话结束后调用：将 turn 写入队列，异步 commit 到 ZhiYi。"""
        if not self._started or not self._client:
            return

        combined = _combine_turn(user_content, assistant_content)
        if not _is_valid_memory_content(combined):
            return

        self._turn_counter += 1
        entry = {
            "content": combined,
            "category": "episodes",
            "metadata": {
                "session_id": session_id or self._session_id,
                "turn": self._turn_counter,
                "platform": self._platform,
            },
        }

        with self._queue_lock:
            self._write_queue.append(entry)

        # 队列超过一定量就 flush，避免积压
        if len(self._write_queue) >= 3:
            self._flush_queue()

    def _flush_queue(self) -> None:
        """将队列中的记忆批量写入 ZhiYi。"""
        with self._queue_lock:
            if not self._write_queue:
                return
            entries = self._write_queue[:]
            self._write_queue.clear()

        # 从 metadata 中读取 agent_id（sync_turn 写入时携带），否则用默认值
        for entry in entries:
            meta = entry.get("metadata", {}) or {}
            self._client.commit(
                content=entry["content"],
                category=entry.get("category", "episodes"),
                metadata=meta,
                agent_id=meta.get("agent_id", "hermes-a06"),
                namespace=meta.get("namespace"),
            )

    # ── Read path ────────────────────────────────────────────────────────────

    def prefetch(self, query: str, *, session_id: str = "") -> str:
        """每次 API 调用前触发：执行语义搜索 + 图谱导航 Obsidian 笔记，返回最相关记忆。"""
        if not self._client or not query or len(query.strip()) < 2:
            return ""

        blocks = ["[ZhiYi Memory — relevant past context]"]

        # 语义搜索
        results = self._client.recall(query.strip(), top_k=3,
                                      agent_id=os.environ.get("ZHIYI_AGENT_ID", "hermes-a06"))
        for r in results:
            score = r.get("score", 0)
            content = r.get("content", "") or r.get("text", "")
            cat = r.get("category", "")
            if content:
                blocks.append(f"  [{score:.2f}][{cat}] {content[:500]}")

        # 图谱导航 + Obsidian 笔记（从 query 提取关键词作为实体）
        notes = self._client.search_notes(query.strip(), max_hops=2, max_notes=3)
        if notes:
            blocks.append("\n[ZhiYi Graph — related Obsidian notes]")
            for n in notes:
                title = n.get("title", "无标题")
                path = n.get("path", "")
                snippet = n.get("snippet", "")[:200]
                score = n.get("score", 0)
                entities = ", ".join(n.get("entities", []))
                blocks.append(f"  [{score:.0f}] {title} ({path})")
                blocks.append(f"      \"{snippet}\"")
                if entities:
                    blocks.append(f"      via entities: {entities}")

        text = "\n".join(blocks)

        # 合并 WebSocket prefetch.push 推送事件
        with self._prefetch_lock:
            if self._prefetch_cache and "[ZhiYi Prefetch" in self._prefetch_cache:
                text += "\n" + self._prefetch_cache

        with self._prefetch_lock:
            self._prefetch_cache = text
        return text

    def queue_prefetch(self, query: str, *, session_id: str = "") -> None:
        """空实现 — prefetch 已是同步的，不需要额外的异步队列。"""
        pass

    # ── Tool interface ────────────────────────────────────────────────────────

    def get_tool_schemas(self) -> List[Dict[str, Any]]:
        return [
            {
                "name": "memory_search",
                "description": "Search ZhiYi semantic memory for relevant past information. Use for: user preferences, earlier decisions, established facts, project context.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "Semantic search query describing what to find. Be specific — the query is encoded with bge-m3 1024-dim and matched against all memories.",
                        },
                        "top_k": {
                            "type": "integer",
                            "description": "Maximum number of results to return (default: 5, max: 20).",
                            "default": 5,
                        },
                    },
                    "required": ["query"],
                },
            },
            {
                "name": "memory_write",
                "description": "Explicitly save an important fact or preference to ZhiYi memory. Use when: user tells you something important, you discover a key fact, user confirms a decision.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "content": {
                            "type": "string",
                            "description": "The memory content to save. Should be a complete, meaningful statement. Be specific — include context.",
                        },
                        "category": {
                            "type": "string",
                            "description": "Category for this memory: 'distilled' (important facts), 'episodes' (conversation turns). Default: 'distilled'.",
                            "default": "distilled",
                        },
                    },
                    "required": ["content"],
                },
            },
            {
                "name": "memory_stats",
                "description": "Get ZhiYi memory statistics: total documents, vector dimensions, index status.",
                "parameters": {
                    "type": "object",
                    "properties": {},
                },
            },
            {
                "name": "memory_feedback",
                "description": "Mark a memory as useful or not-useful. Feedback activates ZhiYi's self-optimization — useful memories get higher priority, not-useful memories trigger quality monitoring and potential deprecation.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "memory_id": {
                            "type": "string",
                            "description": "The memory ID to provide feedback on (from memory_search results).",
                        },
                        "useful": {
                            "type": "boolean",
                            "description": "True = this memory was helpful, False = this memory was not helpful.",
                        },
                        "reason": {
                            "type": "string",
                            "description": "Why this memory was not useful (optional, only used when useful=false).",
                        },
                    },
                    "required": ["memory_id", "useful"],
                },
            },
            {
                "name": "memory_metrics",
                "description": "Get ZhiYi self-optimization metrics: recall hit rate, usefulness rate, gap closure rate, distillation loss, and more. Use this to check memory system health.",
                "parameters": {
                    "type": "object",
                    "properties": {},
                },
            },
            {
                "name": "memory_graph_navigate",
                "description": "Navigate the ZhiYi knowledge graph to understand relationships around an entity. Shows N-hop relationship network — what concepts, people, or projects are connected, and how. Use when: analyzing project structure, understanding who/what is related to a topic, exploring context dependencies.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "entity": {
                            "type": "string",
                            "description": "The central entity to navigate from (e.g. '牧尘', '织忆', '小唯', 'openclaw'). Chinese and English both work.",
                        },
                        "max_hops": {
                            "type": "integer",
                            "description": "Maximum path depth (default: 2). 1-hop = direct neighbors only; 2-hop = friends-of-friends. Larger values return more paths but are slower.",
                            "default": 2,
                        },
                    },
                    "required": ["entity"],
                },
            },
            {
                "name": "memory_graph_stats",
                "description": "Get ZhiYi knowledge graph statistics: total nodes, edges, and graph density. Use to check graph size and health.",
                "parameters": {
                    "type": "object",
                    "properties": {},
                },
            },
        ]

    def handle_tool_call(self, tool_name: str, args: Dict[str, Any], **kwargs) -> str:
        if tool_name == "memory_search":
            return self._tool_memory_search(args.get("query", ""), args.get("top_k", 5))
        elif tool_name == "memory_write":
            return self._tool_memory_write(args.get("content", ""), args.get("category", "distilled"))
        elif tool_name == "memory_feedback":
            return self._tool_memory_feedback(
                args.get("memory_id", ""), args.get("useful", True), args.get("reason", ""))
        elif tool_name == "memory_metrics":
            return self._tool_memory_metrics()
        elif tool_name == "memory_stats":
            return self._tool_memory_stats()
        elif tool_name == "memory_graph_navigate":
            return self._tool_memory_graph_navigate(
                args.get("entity", ""), args.get("max_hops", 2))
        elif tool_name == "memory_graph_stats":
            return self._tool_memory_graph_stats()
        return tool_error(tool_name, "Unknown tool")

    # ── Tool handlers ─────────────────────────────────────────────────────────

    def _tool_memory_search(self, query: str, top_k: int) -> str:
        if not self._client:
            return json.dumps({"success": False, "error": "ZhiYi client not initialized"})
        if top_k > 20:
            top_k = 20
        results = self._client.recall(query, top_k=top_k)
        if not results:
            return json.dumps({"success": True, "results": [], "message": "No relevant memories found"})
        formatted = []
        for r in results:
            content = r.get("content", "") or r.get("text", "")
            formatted.append({
                "content": content[:1000],
                "score": round(r.get("score", 0), 4),
                "category": r.get("category", ""),
                "quality_score": round(r.get("quality_score", 0), 4),
                "v_value": round(r.get("v_value", 0), 4),
            })
        return json.dumps({"success": True, "count": len(formatted), "results": formatted})

    def _tool_memory_write(self, content: str, category: str) -> str:
        if not self._client:
            return json.dumps({"success": False, "error": "ZhiYi client not initialized"})
        if not _is_valid_memory_content(content):
            return json.dumps({"success": False, "error": "Content too short or invalid"})
        commit_id = self._client.commit(
            content=content, category=category,
            agent_id=os.environ.get("ZHIYI_AGENT_ID", "hermes-a06"),
        )
        if commit_id:
            return json.dumps({"success": True, "commit_id": commit_id})
        return json.dumps({"success": False, "error": "Commit failed — check ZhiYi server logs"})

    def _tool_memory_stats(self) -> str:
        if not self._client:
            return json.dumps({"success": False, "error": "ZhiYi client not initialized"})
        stats = self._client.stats()
        if not stats:
            return json.dumps({"success": False, "error": "Could not reach ZhiYi /stats endpoint"})
        return json.dumps({"success": True, "stats": stats})

    def _tool_memory_feedback(self, memory_id: str, useful: bool, reason: str = "") -> str:
        if not self._client:
            return json.dumps({"success": False, "error": "ZhiYi client not initialized"})
        if not memory_id:
            return json.dumps({"success": False, "error": "memory_id required"})
        if useful:
            ok = self._client.mark_useful(memory_id)
        else:
            ok = self._client.mark_not_useful(memory_id, reason)
        if ok:
            return json.dumps({"success": True, "memory_id": memory_id, "useful": useful,
                               "message": "Feedback recorded — VProp + Dashboard updated"})
        return json.dumps({"success": False, "error": "ZhiYi feedback API failed"})

    def _tool_memory_metrics(self) -> str:
        if not self._client:
            return json.dumps({"success": False, "error": "ZhiYi client not initialized"})
        metrics = self._client.metrics()
        if not metrics:
            return json.dumps({"success": False, "error": "Could not reach ZhiYi /metrics/self"})
        return json.dumps({"success": True, "metrics": metrics})

    def _tool_memory_graph_navigate(self, entity: str, max_hops: int) -> str:
        """图谱导航：查询实体的 N 跳关系网络。"""
        if not self._client:
            return json.dumps({"success": False, "error": "ZhiYi client not initialized"})
        if not entity or len(entity.strip()) < 1:
            return json.dumps({"success": False, "error": "entity is required"})
        if max_hops < 1:
            max_hops = 1
        if max_hops > 5:
            max_hops = 5

        result = self._client.graph_navigate(entity.strip(), max_hops=max_hops)
        if not result:
            return json.dumps({"success": False, "error": "Graph navigate failed — check ZhiYi server logs"})

        # 格式化路径输出，保留关键信息
        paths = result.get("paths", [])
        formatted_paths = []
        for p in paths[:50]:  # 最多返回50条路径，避免太长
            formatted_paths.append({
                "from": p.get("from", ""),
                "relation": p.get("relation", ""),
                "to": p.get("to", ""),
                "hop": p.get("hop", 1),
                "weight": round(p.get("weight", 0), 3),
            })

        return json.dumps({
            "success": True,
            "entity": result.get("entity", entity),
            "bidirectional": result.get("bidirectional", False),
            "count": result.get("count", len(formatted_paths)),
            "paths": formatted_paths,
            "message": f"Found {result.get('count', 0)} paths within {max_hops}-hop network",
        })

    def _tool_memory_graph_stats(self) -> str:
        """获取图谱统计：节点数、边数、密度。"""
        if not self._client:
            return json.dumps({"success": False, "error": "ZhiYi client not initialized"})
        stats = self._client.graph_stats()
        if not stats:
            return json.dumps({"success": False, "error": "Could not reach ZhiYi /graph/stats"})
        return json.dumps({
            "success": True,
            "stats": stats,
            "message": f"Graph has {stats.get('node_count', 0)} nodes and {stats.get('edge_count', 0)} edges (density: {stats.get('density', 0):.4f})",
        })

    # ── Session management ────────────────────────────────────────────────────

    def on_session_switch(self, new_session_id: str, *, parent_session_id: str = "", reset: bool = False, **kwargs) -> None:
        self._session_id = new_session_id
        self._turn_counter = 0
        if reset:
            self._flush_queue()

    def on_turn_start(self, turn_number: int, message: str, **kwargs) -> None:
        """每轮开始时 flush 待写入的记忆，保证顺序。"""
        self._flush_queue()

    # ── System prompt ────────────────────────────────────────────────────────

    def system_prompt_block(self) -> str:
        return (
            "\\n[ZhiYi Memory] You have access to ZhiYi MemoryWeave — a semantic memory system with knowledge graph.\\n"
            "Your prefetch automatically retrieves: (1) relevant semantic memories, (2) Obsidian notes related via graph navigation.\\n"
            "Tools available:\\n"
            "  memory_search — semantic search for relevant past information\\n"
            "  memory_write — save an important fact or preference\\n"
            "  memory_feedback — mark memories as useful/not-useful (drives self-optimization)\\n"
            "  memory_metrics — check memory system health (recall hit rate, gap closures, etc.)\\n"
            "  memory_stats — get memory statistics (total docs, vector dimensions)\\n"
            "  memory_graph_navigate — navigate the knowledge graph to understand entity relationships (N-hop network)\\n"
            "  memory_graph_stats — get knowledge graph statistics (nodes, edges, density)\\n"
        )


# ── Plugin entry point ──────────────────────────────────────────────────────


def register(ctx) -> None:
    """Called by Hermes plugin system to register this memory provider."""
    ctx.register_memory_provider(HermesZhiYiMemoryProvider())