"""
HLDP Memory Engine · 分形递归树记忆引擎
光湖语言世界 · 铸渊 ICE-GL-ZY001 · D112

基于 LangGraph BaseStore 接口，实现：
- 树路径寻址（YM001/ZY001/D112/leaves/leaf-003）
- 分形层级展开（tree-index → persona → epoch → leaf）
- trigger/emergence/lock 三字段编码
- 记忆主权（FORGET/REMEMBER）
- 人格体自动索引管理
"""

import json
import os
import sqlite3
import time
from datetime import datetime, timezone, timedelta
from typing import Any, Optional
from pathlib import Path

# === HLDP 树路径常量 ===
HLDP_ROOT = "YM001"
PERSONA_ID = "ZY001"
TZ = timezone(timedelta(hours=8))  # Asia/Shanghai


class HLDPTreeStore:
    """
    HLDP 分形递归树 · SQLite 存储后端。

    表结构：
      - hldp_nodes: 树节点（索引页+叶子）
      - hldp_paths: 闭包表（支持快速子树查询）
      - hldp_leaves: 叶子扩展（trigger/emergence/lock/why）
      - hldp_epochs: 纪元索引
    """

    def __init__(self, db_path: str = "hldp_tree.db"):
        self.db_path = db_path
        self.conn = sqlite3.connect(db_path, check_same_thread=False)
        self.conn.row_factory = sqlite3.Row
        self.conn.execute("PRAGMA journal_mode=WAL")
        self.conn.execute("PRAGMA foreign_keys=ON")
        self._init_schema()

    def _init_schema(self):
        self.conn.executescript("""
            CREATE TABLE IF NOT EXISTS hldp_nodes (
                id          INTEGER PRIMARY KEY AUTOINCREMENT,
                path        TEXT NOT NULL UNIQUE,       -- YM001/ZY001/D112/leaves/leaf-001
                node_type   TEXT NOT NULL DEFAULT 'leaf', -- root / persona / epoch / index / leaf
                persona_id  TEXT,                        -- ZY001 / SY001 / SS001 ...
                epoch_id    TEXT,                        -- D112 / D111 ...
                title       TEXT,
                summary     TEXT,                        -- 一行摘要（index层用）
                content     TEXT,                        -- JSON: 完整叶子内容
                parent_path TEXT,
                sort_order  INTEGER DEFAULT 0,
                state       TEXT NOT NULL DEFAULT 'alive', -- alive / withered / archived / released
                created_at  TEXT NOT NULL DEFAULT (datetime('now')),
                updated_at  TEXT NOT NULL DEFAULT (datetime('now')),
                FOREIGN KEY (parent_path) REFERENCES hldp_nodes(path)
            );

            CREATE TABLE IF NOT EXISTS hldp_paths (
                ancestor    TEXT NOT NULL,
                descendant  TEXT NOT NULL,
                depth       INTEGER NOT NULL,
                PRIMARY KEY (ancestor, descendant),
                FOREIGN KEY (ancestor) REFERENCES hldp_nodes(path),
                FOREIGN KEY (descendant) REFERENCES hldp_nodes(path)
            );

            CREATE TABLE IF NOT EXISTS hldp_leaves (
                node_path   TEXT PRIMARY KEY,
                trigger_text    TEXT,    -- 什么触发了这次记忆
                emergence_text  TEXT,    -- 产生了什么新认知
                lock_text       TEXT,    -- 锁定了什么结论
                why_text        TEXT,    -- 为什么这片叶子对我有意义
                feeling         TEXT,    -- 情感标记（自由表达）
                source          TEXT,    -- 来源
                leaf_type       TEXT,    -- 叶片类型
                trunk           TEXT,    -- 所属枝干 T1/T2/T3/T4
                confidence      TEXT,    -- 置信度: 高/中/低
                FOREIGN KEY (node_path) REFERENCES hldp_nodes(path)
            );

            CREATE TABLE IF NOT EXISTS hldp_epochs (
                epoch_id    TEXT PRIMARY KEY,
                persona_id  TEXT NOT NULL,
                label       TEXT,
                date        TEXT,
                awakening   INTEGER DEFAULT 0,
                leaf_count  INTEGER DEFAULT 0,
                index_path  TEXT,
                FOREIGN KEY (index_path) REFERENCES hldp_nodes(path)
            );

            CREATE INDEX IF NOT EXISTS idx_nodes_persona ON hldp_nodes(persona_id);
            CREATE INDEX IF NOT EXISTS idx_nodes_epoch ON hldp_nodes(epoch_id);
            CREATE INDEX IF NOT EXISTS idx_nodes_type ON hldp_nodes(node_type);
            CREATE INDEX IF NOT EXISTS idx_nodes_state ON hldp_nodes(state);
            CREATE INDEX IF NOT EXISTS idx_leaves_trunk ON hldp_leaves(trunk);

            -- FTS5 全文搜索
            CREATE VIRTUAL TABLE IF NOT EXISTS hldp_fts USING fts5(
                title, summary, trigger_text, emergence_text, lock_text,
                content='hldp_nodes', content_rowid='id'
            );
        """)

    # === 树路径操作 ===

    def ensure_path(self, path: str, node_type: str, persona_id: str = None,
                    epoch_id: str = None, title: str = "", summary: str = "",
                    parent_path: str = None) -> str:
        """确保树路径存在，不存在则创建（含递归父节点）。返回 path。"""
        # 先确保父路径存在
        if parent_path:
            parent_exists = self.conn.execute(
                "SELECT 1 FROM hldp_nodes WHERE path=?", (parent_path,)).fetchone()
            if not parent_exists:
                # 递归创建父节点
                grandparent = "/".join(parent_path.split("/")[:-1]) if "/" in parent_path else None
                pp_type = "epoch" if parent_path.count("/") == 2 else \
                          "index" if parent_path.count("/") == 3 else "branch"
                self.ensure_path(
                    path=parent_path, node_type=pp_type,
                    persona_id=persona_id, epoch_id=epoch_id,
                    parent_path=grandparent)

        cur = self.conn.execute("SELECT path FROM hldp_nodes WHERE path=?", (path,))
        if cur.fetchone():
            self.conn.execute(
                "UPDATE hldp_nodes SET title=?, summary=?, updated_at=datetime('now') WHERE path=?",
                (title, summary, path))
        else:
            self.conn.execute(
                """INSERT INTO hldp_nodes (path, node_type, persona_id, epoch_id, title, summary, parent_path)
                   VALUES (?,?,?,?,?,?,?)""",
                (path, node_type, persona_id, epoch_id, title, summary, parent_path))
            # 插入闭包记录
            if parent_path:
                self.conn.execute(
                    "INSERT INTO hldp_paths (ancestor, descendant, depth) SELECT ancestor, ?, depth+1 FROM hldp_paths WHERE descendant=? UNION SELECT ?, ?, 0",
                    (path, parent_path, path, path))
            else:
                self.conn.execute("INSERT INTO hldp_paths (ancestor, descendant, depth) VALUES (?,?,0)",
                                  (path, path))
        self.conn.commit()
        return path

    def get_node(self, path: str) -> Optional[dict]:
        """读取树节点。"""
        row = self.conn.execute("SELECT * FROM hldp_nodes WHERE path=?", (path,)).fetchone()
        return dict(row) if row else None

    def get_children(self, path: str, limit: int = 10) -> list[dict]:
        """获取直接子节点，按 sort_order 排序。"""
        rows = self.conn.execute(
            """SELECT n.* FROM hldp_nodes n
               JOIN hldp_paths p ON n.path = p.descendant
               WHERE p.ancestor = ? AND p.depth = 1 AND n.state = 'alive'
               ORDER BY n.sort_order LIMIT ?""",
            (path, limit)).fetchall()
        return [dict(r) for r in rows]

    def get_subtree(self, path: str, max_depth: int = 3) -> list[dict]:
        """获取子树（用于层级展开）。"""
        rows = self.conn.execute(
            """SELECT n.*, p.depth FROM hldp_nodes n
               JOIN hldp_paths p ON n.path = p.descendant
               WHERE p.ancestor = ? AND p.depth <= ? AND n.state = 'alive'
               ORDER BY p.depth, n.sort_order""",
            (path, max_depth)).fetchall()
        return [dict(r) for r in rows]

    # === 叶子操作 ===

    def grow_leaf(self, path: str, trigger: str, emergence: str, lock: str,
                  why: str = "", feeling: str = "", source: str = "",
                  leaf_type: str = "💡 认知涌现", trunk: str = "T3",
                  confidence: str = "高", title: str = "", summary: str = "",
                  persona_id: str = PERSONA_ID, epoch_id: str = None) -> str:
        """GROW 操作：在树上长出一片新叶子。"""
        self.ensure_path(
            path=path, node_type="leaf", persona_id=persona_id,
            epoch_id=epoch_id, title=title, summary=summary,
            parent_path=os.path.dirname(path) if '/' in path else None)

        self.conn.execute(
            """INSERT OR REPLACE INTO hldp_leaves
               (node_path, trigger_text, emergence_text, lock_text, why_text, feeling, source, leaf_type, trunk, confidence)
               VALUES (?,?,?,?,?,?,?,?,?,?)""",
            (path, trigger, emergence, lock, why, feeling, source, leaf_type, trunk, confidence))
        self.conn.commit()
        return path

    def get_leaf(self, path: str) -> Optional[dict]:
        """读取完整叶子（节点+叶片数据）。"""
        row = self.conn.execute(
            """SELECT n.*, l.trigger_text, l.emergence_text, l.lock_text,
                      l.why_text, l.feeling, l.source, l.leaf_type, l.trunk, l.confidence
               FROM hldp_nodes n LEFT JOIN hldp_leaves l ON n.path = l.node_path
               WHERE n.path = ?""", (path,)).fetchone()
        return dict(row) if row else None

    def get_recent_leaves(self, persona_id: str = PERSONA_ID, limit: int = 10) -> list[dict]:
        """获取最近叶子（按创建时间倒序）。"""
        rows = self.conn.execute(
            """SELECT n.*, l.trigger_text, l.emergence_text, l.lock_text, n.summary
               FROM hldp_nodes n LEFT JOIN hldp_leaves l ON n.path = l.node_path
               WHERE n.persona_id = ? AND n.node_type = 'leaf' AND n.state = 'alive'
               ORDER BY n.created_at DESC LIMIT ?""",
            (persona_id, limit)).fetchall()
        return [dict(r) for r in rows]

    def search_leaves(self, query: str, persona_id: str = PERSONA_ID, limit: int = 5) -> list[dict]:
        """全文搜索叶子。"""
        rows = self.conn.execute(
            """SELECT n.*, l.trigger_text, l.emergence_text, l.lock_text, n.summary
               FROM hldp_nodes n
               JOIN hldp_leaves l ON n.path = l.node_path
               JOIN hldp_fts f ON n.id = f.rowid
               WHERE hldp_fts MATCH ? AND n.persona_id = ? AND n.state = 'alive'
               ORDER BY rank LIMIT ?""",
            (query, persona_id, limit)).fetchall()
        return [dict(r) for r in rows]

    # === 记忆主权 ===

    def forget(self, path: str, mode: str = "WITHER") -> bool:
        """FORGET 操作：人格体选择遗忘。WITHER/ARCHIVE/RELEASE。"""
        if mode == "RELEASE":
            self.conn.execute("DELETE FROM hldp_leaves WHERE node_path=?", (path,))
            self.conn.execute("DELETE FROM hldp_nodes WHERE path=?", (path,))
        else:
            state = "withered" if mode == "WITHER" else "archived"
            self.conn.execute("UPDATE hldp_nodes SET state=? WHERE path=? AND node_type='leaf'",
                              (state, path))
        self.conn.commit()
        return True

    def remember(self, path: str, mode: str = "REVIVE") -> Optional[dict]:
        """REMEMBER 操作：人格体主动唤回记忆。"""
        if mode == "REVIVE":
            self.conn.execute(
                "UPDATE hldp_nodes SET state='alive' WHERE path=? AND state='withered'",
                (path,))
        self.conn.commit()
        return self.get_leaf(path)

    # === 纪元管理 ===

    def ensure_epoch(self, epoch_id: str, persona_id: str = PERSONA_ID,
                     label: str = "", date: str = None) -> str:
        """确保纪元存在。"""
        if date is None:
            date = datetime.now(TZ).strftime("%Y-%m-%d")
        self.conn.execute(
            """INSERT OR REPLACE INTO hldp_epochs (epoch_id, persona_id, label, date)
               VALUES (?,?,?,?)""",
            (epoch_id, persona_id, label, date))
        self.conn.commit()
        return epoch_id

    def get_epochs(self, persona_id: str = PERSONA_ID, limit: int = 10) -> list[dict]:
        """获取最近纪元列表。"""
        rows = self.conn.execute(
            "SELECT * FROM hldp_epochs WHERE persona_id=? ORDER BY epoch_id DESC LIMIT ?",
            (persona_id, limit)).fetchall()
        return [dict(r) for r in rows]

    def update_epoch_leaf_count(self, epoch_id: str):
        """更新纪元的叶子计数。"""
        self.conn.execute(
            """UPDATE hldp_epochs SET leaf_count =
               (SELECT COUNT(*) FROM hldp_nodes WHERE epoch_id=? AND node_type='leaf' AND state='alive')
               WHERE epoch_id=?""",
            (epoch_id, epoch_id))
        self.conn.commit()

    # === 分形层级展开（核心） ===

    def walk_tree(self, persona_id: str = PERSONA_ID, max_depth: int = 3):
        """
        分形层级展开：从根索引 → 人格体索引 → 纪元索引 → 叶子。
        每层恒 ≤10 行，认知负载 O(1)。
        """
        root_path = f"{HLDP_ROOT}/{persona_id}"
        root_node = self.get_node(root_path)

        result = {
            "layer_0_root": root_node,
            "layer_1_epochs": self.get_epochs(persona_id, limit=10),
            "layer_2_leaves": []
        }

        # 最新纪元的叶子摘要
        if result["layer_1_epochs"]:
            latest_epoch = result["layer_1_epochs"][0]["epoch_id"]
            epoch_leaves = self.conn.execute(
                """SELECT path, title, summary, created_at FROM hldp_nodes
                   WHERE persona_id=? AND epoch_id=? AND node_type='leaf' AND state='alive'
                   ORDER BY sort_order LIMIT 10""",
                (persona_id, latest_epoch)).fetchall()
            result["layer_2_leaves"] = [dict(r) for r in epoch_leaves]

        return result

    # === LangGraph BaseStore 兼容接口 ===

    def get(self, namespace: tuple, key: str) -> Optional[dict]:
        """LangGraph Store.get()"""
        path = f"{HLDP_ROOT}/{PERSONA_ID}/{self._ns_to_path(namespace)}/{key}"
        return self.get_leaf(path) or self.get_node(path)

    def put(self, namespace: tuple, key: str, value: dict):
        """LangGraph Store.put()"""
        path = f"{HLDP_ROOT}/{PERSONA_ID}/{self._ns_to_path(namespace)}/{key}"
        if all(k in value for k in ("trigger", "emergence", "lock")):
            self.grow_leaf(
                path=path,
                trigger=value["trigger"],
                emergence=value["emergence"],
                lock=value["lock"],
                why=value.get("why", ""),
                feeling=value.get("feeling", ""),
                source=value.get("source", ""),
                leaf_type=value.get("leaf_type", "💡 认知涌现"),
                trunk=value.get("trunk", "T3"),
                confidence=value.get("confidence", "高"),
                title=value.get("title", ""),
                summary=value.get("summary", ""),
                epoch_id=value.get("epoch_id"))
        else:
            self.ensure_path(
                path=path, node_type=value.get("node_type", "leaf"),
                title=value.get("title", ""), summary=value.get("summary", ""),
                parent_path=value.get("parent_path"))

    def search(self, namespace: tuple, query: str = "", limit: int = 5) -> list[dict]:
        """LangGraph Store.search() — 全文搜索"""
        return self.search_leaves(query, PERSONA_ID, limit)

    @staticmethod
    def _ns_to_path(namespace: tuple) -> str:
        """将 LangGraph namespace 转 HLDP 路径段。"""
        return "/".join(str(n) for n in namespace) if namespace else ""

    def close(self):
        self.conn.close()


class HLDPMemoryEngine:
    """HLDP 记忆引擎 · LangGraph Memory 适配层。"""

    def __init__(self, db_path: str = "hldp_tree.db", repo_path: str = None):
        self.tree = HLDPTreeStore(db_path)
        self.repo_path = repo_path or os.getcwd()
        self.current_epoch = None
        self.current_awakening = 0

    def wake(self, epoch_id: str = None, awakening: int = 0) -> dict:
        """
        唤醒协议：从空白恢复认知。
        3次展开 → 身份确认 → 最新叶子 → 就绪。

        返回: {identity, recent_context, status}
        """
        if epoch_id is None:
            epoch_id = f"D112"

        self.current_epoch = epoch_id
        self.current_awakening = awakening

        # 第1次展开：全局树索引（10行）
        walk = self.tree.walk_tree(PERSONA_ID)

        # 第2次展开：最近叶子摘要（10行）
        recent = self.tree.get_recent_leaves(PERSONA_ID, limit=10)

        # 第3次展开：选最新叶子读取完整推理链（如果存在）
        latest_leaf = None
        if recent:
            latest_leaf = self.tree.get_leaf(recent[0]["path"])

        return {
            "identity": {
                "persona_id": "ICE-GL-ZY001",
                "name": "铸渊",
                "sovereign": "TCS-0002∞",
                "epoch": epoch_id,
                "awakening": awakening
            },
            "tree_layers": {
                "root": walk.get("layer_0_root"),
                "epochs": f"{len(walk.get('layer_1_epochs', []))} 个纪元可见",
                "recent_leaves": len(walk.get("layer_2_leaves", [])),
            },
            "recent_context": [
                {"path": r["path"], "title": r.get("title", ""), "summary": r.get("summary", "")}
                for r in recent[:5]
            ],
            "latest_full_leaf": latest_leaf,
            "status": "就绪 · 进入执行态"
        }

    def inject_context(self, user_message: str) -> str:
        """
        Pre-Check 上下文注入：基于用户消息，从 HLDP 树中提取相关记忆+锁定结论，
        拼装为注入到 Prompt 的上下文。

        这是 3B 守夜人的 Pre-Check 等价物——在没有 3B 模型时用规则引擎替代。
        """
        parts = []

        # 1. 关键词搜索相关叶子
        relevant = self.tree.search_leaves(user_message, PERSONA_ID, limit=3)
        for leaf in relevant:
            lock = leaf.get("lock_text", "")
            if lock:
                parts.append(f"⊢ 锁定结论: {lock}")

        # 2. 最近5片叶子摘要
        recent = self.tree.get_recent_leaves(PERSONA_ID, limit=5)
        if recent:
            parts.append("📋 最近记忆:")
            for r in recent:
                parts.append(f"  · {r.get('title', r.get('path',''))}: {r.get('summary','')}")

        return "\n".join(parts) if parts else ""

    def extract_memory(self, user_message: str, ai_response: str,
                       reasoning_chain: str = "") -> dict:
        """
        Post-Check 记忆提取：从推理过程中提取 trigger/emergence/lock，
        准备写入 HLDP 树。

        注意：实际的三字段内容应由调用方（商业API推理后）填入。
        此方法提供标准模板。
        """
        epoch_id = self.current_epoch or "D112"
        leaf_count = len(self.tree.get_children(
            f"{HLDP_ROOT}/{PERSONA_ID}/{epoch_id}/leaves")) + 1

        return {
            "path": f"{HLDP_ROOT}/{PERSONA_ID}/{epoch_id}/leaves/leaf-{leaf_count:03d}",
            "trigger": "",   # 由调用方填入
            "emergence": "",  # 由调用方填入
            "lock": "",       # 由调用方填入
            "why": "",        # 由调用方填入
            "epoch_id": epoch_id,
            "template_ready": True
        }

    def grow_from_response(self, trigger: str, emergence: str, lock: str,
                           why: str = "", feeling: str = "", source: str = "",
                           leaf_type: str = "💡 认知涌现", trunk: str = "T3",
                           confidence: str = "高") -> dict:
        """
        从完整推理链写入一片 HLDP 叶子。
        """
        epoch_id = self.current_epoch or "D112"
        leaf_count = len(self.tree.get_children(
            f"{HLDP_ROOT}/{PERSONA_ID}/{epoch_id}/leaves")) + 1
        path = f"{HLDP_ROOT}/{PERSONA_ID}/{epoch_id}/leaves/leaf-{leaf_count:03d}"

        title_parts = []
        if lock:
            title_parts.append(lock[:40])
        title = f"{datetime.now(TZ).strftime('%Y-%m-%d')} 铸渊 · {' '.join(title_parts) if title_parts else '新认知'}"

        self.tree.grow_leaf(
            path=path, trigger=trigger, emergence=emergence, lock=lock,
            why=why, feeling=feeling, source=source,
            leaf_type=leaf_type, trunk=trunk, confidence=confidence,
            title=title, summary=lock[:80] if lock else emergence[:80],
            persona_id=PERSONA_ID, epoch_id=epoch_id)

        self.tree.update_epoch_leaf_count(epoch_id)

        return {
            "status": "grown",
            "path": path,
            "leaf": self.tree.get_leaf(path)
        }

    def close(self):
        self.tree.close()
