import sqlite3
import os
import sys
import re
import json
import uuid
import time
import shutil
from typing import List, Dict, Optional

# 🝙 P0 fix: DB path based on the file's location (no longer depends on the process's working directory)
# 🌟 打包（PyInstaller onedir）时：优先与源码树共用 data/chat_history.db（历史不丢失）；
#    若 exe 被复制到源码树之外 → 回退为 exe 旁的 data/chat_history.db
if getattr(sys, "frozen", False):
    _exe_dir = os.path.dirname(os.path.abspath(sys.executable))
    _shared = os.path.abspath(os.path.join(_exe_dir, "..", "..", "data", "chat_history.db"))
    if os.path.isdir(os.path.dirname(_shared)):
        _DEFAULT_DB = _shared
    else:
        _DEFAULT_DB = os.path.join(_exe_dir, "data", "chat_history.db")
else:
    _DEFAULT_DB = os.path.join(
        os.path.dirname(os.path.abspath(__file__)), "..", "data", "chat_history.db")


class _AutoCloseConn:
    """🝙 P0 fix: wraps sqlite3.Connection.

    sqlite3.Connection's with only handles commit/rollback and does NOT close.
    All 19 call sites use `with self.get_connection() as conn:`, so we auto-close on with exit.
    """

    def __init__(self, conn):
        self._conn = conn

    def __getattr__(self, name):
        return getattr(self._conn, name)

    def __enter__(self):
        self._conn.__enter__()
        return self

    def __exit__(self, exc_type, exc, tb):
        try:
            return self._conn.__exit__(exc_type, exc, tb)
        finally:
            try:
                self._conn.close()
            except Exception:
                pass


class DBManager:
    def __init__(self, db_path: str = None):
        # 🝙 Default path is anchored to the core/ directory, not affected by launch CWD
        self.db_path = db_path or os.path.abspath(_DEFAULT_DB)

        # 🆕 附件/媒体文件根目录（copy_session 深拷贝磁盘文件用）。默认=项目根；测试可覆盖。
        self.files_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

        self.is_first_run = not os.path.exists(self.db_path)
        d = os.path.dirname(self.db_path)
        if d:
            os.makedirs(d, exist_ok=True)

        self._init_db()

    def get_connection(self):
        conn = sqlite3.connect(self.db_path, check_same_thread=False)
        conn.row_factory = sqlite3.Row
        conn.execute("PRAGMA foreign_keys = ON")
        return _AutoCloseConn(conn)

    def _init_db(self):
        """初始化表结构并注入默认数据"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            
            # 1. 创建 sessions 表 (新增 current_leaf_msg_id)
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS sessions (
                    id TEXT PRIMARY KEY,
                    title TEXT,
                    created_at INTEGER,
                    updated_at INTEGER,
                    has_messages BOOLEAN DEFAULT 0,
                    sort_order INTEGER DEFAULT 0,
                    is_starred BOOLEAN DEFAULT 0,
                    current_leaf_msg_id TEXT
                )
            """)
            
            # 2. 创建 messages 表 (新增 parent_id)
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS messages (
                    id TEXT PRIMARY KEY,
                    session_id TEXT,
                    role TEXT,
                    content TEXT,
                    reasoning TEXT,
                    is_ignored BOOLEAN,
                    created_at INTEGER,
                    attachment_metadata TEXT,
                    parent_id TEXT,
                    timeline TEXT,
                    usage TEXT,
                    stop_reason TEXT,
                    error_message TEXT,
                    FOREIGN KEY(session_id) REFERENCES sessions(id) ON DELETE CASCADE
                )
            """)

            # ==========================================
            # 🌟 核心性能优化：为高频查询的字段建立索引
            # ==========================================
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_session_id ON messages(session_id)")
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_parent_id ON messages(parent_id)")

            # 3. 🌟 自动化热升级：检测并兼容旧数据库
            self._upgrade_schema(cursor)
            
            # 4. 检查是否需要插入初始默认对话
            cursor.execute("SELECT COUNT(*) FROM sessions")
            if cursor.fetchone()[0] == 0:
                self._seed_default_chat(cursor)
                
            conn.commit()


    def _upgrade_schema(self, cursor):
        """检测缺少的新字段并自动补齐，如果是刚升级，则自动将旧线性数据串联成链表"""
        upgraded = False
        
        # 兼容 sessions 字段
        for col in ["has_messages", "sort_order", "is_starred"]:
            try:
                cursor.execute(f"SELECT {col} FROM sessions LIMIT 1")
            except sqlite3.OperationalError:
                if col == "has_messages": cursor.execute("ALTER TABLE sessions ADD COLUMN has_messages BOOLEAN DEFAULT 0")
                if col == "sort_order": cursor.execute("ALTER TABLE sessions ADD COLUMN sort_order INTEGER DEFAULT 0")
                if col == "is_starred": cursor.execute("ALTER TABLE sessions ADD COLUMN is_starred BOOLEAN DEFAULT 0")

        # 🌟 核心：兼容链表树架构
        try:
            # 兼容 sessions.mode 列（chat/worker 模式锁定，NULL=未发送过）
            try:
                cursor.execute("SELECT mode FROM sessions LIMIT 1")
            except Exception:
                cursor.execute("ALTER TABLE sessions ADD COLUMN mode TEXT")

            cursor.execute("SELECT current_leaf_msg_id FROM sessions LIMIT 1")
        except sqlite3.OperationalError:
            print("[DB]: 正在升级 sessions 表 (添加 current_leaf_msg_id)...")
            cursor.execute("ALTER TABLE sessions ADD COLUMN current_leaf_msg_id TEXT")
            upgraded = True

        try:
            cursor.execute("SELECT attachment_metadata, parent_id FROM messages LIMIT 1")
        except sqlite3.OperationalError:
            print("[DB]: 正在升级 messages 表 (添加 attachment_metadata, parent_id)...")
            try: cursor.execute("ALTER TABLE messages ADD COLUMN attachment_metadata TEXT")
            except: pass
            try: cursor.execute("ALTER TABLE messages ADD COLUMN parent_id TEXT")
            except: pass
            upgraded = True

        # 🌟 messages.timeline 列（agent 时间线 JSON: 思考/文本/工具 按事件顺序）
        try:
            cursor.execute("SELECT timeline FROM messages LIMIT 1")
        except sqlite3.OperationalError:
            print("[DB]: 正在升级 messages 表 (添加 timeline)...")
            try: cursor.execute("ALTER TABLE messages ADD COLUMN timeline TEXT")
            except: pass
            upgraded = True

        # 🆕 P1: messages.usage 列（assistant 回复的精确 usage JSON，
        #        供显示/压缩估算做 usage 锚定，对照 pi 内存态 usage 回放）
        try:
            cursor.execute("SELECT usage FROM messages LIMIT 1")
        except sqlite3.OperationalError:
            print("[DB]: 正在升级 messages 表 (添加 usage)...")
            try: cursor.execute("ALTER TABLE messages ADD COLUMN usage TEXT")
            except: pass
            upgraded = True

        # 🆕 失败轮次持久化（对照 pi: message_end 无条件入库）：
        #    messages.stop_reason / error_message —— 区分「正常行 / 出错行」，
        #    供 UI 展示与 build_api_context 回放决策（NULL = 正常）
        #
        # ⚠️⚠️ 绝不能置 upgraded=True：该标志会触发下方的「旧数据链表化重构」，
        #    把用户的**树状分支拍平成线性链**（数据破坏）！
        #    纯追加列对本迁移自身而言是安全的，与旧库结构修复无关。
        try:
            cursor.execute("SELECT stop_reason, error_message FROM messages LIMIT 1")
        except sqlite3.OperationalError:
            print("[DB]: 正在升级 messages 表 (添加 stop_reason, error_message)...")
            try: cursor.execute("ALTER TABLE messages ADD COLUMN stop_reason TEXT")
            except: pass
            try: cursor.execute("ALTER TABLE messages ADD COLUMN error_message TEXT")
            except: pass

        # 如果刚才执行了树状结构升级，立即对旧数据进行“时间线串联”修复
        if upgraded:
            print("[DB]: 🚀 正在执行旧数据链表化重构...")
            cursor.execute("SELECT id FROM sessions")
            sessions = cursor.fetchall()
            for s in sessions:
                sid = s['id']
                cursor.execute("SELECT id FROM messages WHERE session_id = ? ORDER BY created_at ASC", (sid,))
                msgs = cursor.fetchall()
                if not msgs: continue
                
                # 遍历消息，将后一条的 parent_id 指向上一条
                prev_id = None
                for m in msgs:
                    mid = m['id']
                    if prev_id:
                        cursor.execute("UPDATE messages SET parent_id = ? WHERE id = ?", (prev_id, mid))
                    prev_id = mid
                
                # 最后一个 msg_id 就是这棵树的末端叶子节点
                cursor.execute("UPDATE sessions SET current_leaf_msg_id = ? WHERE id = ?", (prev_id, sid))
            print("[DB]: ✅ 数据结构完美升级为链表树！")


    def _seed_default_chat(self, cursor):
        session_id = f"sess_{uuid.uuid4().hex[:12]}"
        now = int(time.time())
        
        cursor.execute(
            "INSERT INTO sessions (id, title, created_at, updated_at, has_messages) VALUES (?, ?, ?, ?, ?)",
            (session_id, "让我们从这里开始吧", now, now, 1)
        )
        
        sys_id = f"msg_sys_init"
        cursor.execute("""
            INSERT INTO messages (id, session_id, role, content, reasoning, is_ignored, created_at, parent_id)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?)
        """, (sys_id, session_id, "system", "你是一个优秀的助手！", "", 0, now, None))
        
        default_messages = [
            {"role": "user", "content": "你好呀，你是谁？"},
            {"role": "assistant", "content": "嗨！我是你的 AI 助手 ✨ ..."},
            {"role": "user", "content": "那你到底能帮我做什么？"},
            {"role": "assistant", "content": "简单来说，能打字问的我都聊..."}
        ]
        
        prev_id = sys_id
        for msg in default_messages:
            msg_id = f"msg_{uuid.uuid4().hex[:16]}"
            now += 1 
            cursor.execute("""
                INSERT INTO messages (id, session_id, role, content, reasoning, is_ignored, created_at, parent_id)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            """, (msg_id, session_id, msg["role"], msg["content"], "", 0, now, prev_id))
            prev_id = msg_id
            
        # 设置默认会话的叶子节点
        cursor.execute("UPDATE sessions SET current_leaf_msg_id = ? WHERE id = ?", (prev_id, session_id))

    # ==================== 会话 (Session) 操作 ====================
    def get_session_mode(self, session_id: str) -> Optional[str]:
        """读取会话锁定的模式（chat/worker），未发送过返回 None"""
        with self.get_connection() as conn:
            row = conn.execute("SELECT mode FROM sessions WHERE id = ?",
                               (session_id,)).fetchone()
            return row["mode"] if row else None

    def set_session_mode(self, session_id: str, mode: str):
        """锁定会话模式（首条消息发送时调用，之后不可变）"""
        with self.get_connection() as conn:
            conn.execute("UPDATE sessions SET mode = ? WHERE id = ?",
                         (mode, session_id))

    def get_all_sessions(self) -> List[Dict]:
        with self.get_connection() as conn:
            return [dict(row) for row in conn.execute("SELECT * FROM sessions ORDER BY is_starred DESC, sort_order ASC, updated_at DESC").fetchall()]

    def create_session(self, title: str = "新对话") -> Dict:
        session_id = f"sess_{uuid.uuid4().hex[:12]}"
        now = int(time.time())
        sys_msg_id = f"msg_{uuid.uuid4().hex[:16]}"
        
        with self.get_connection() as conn:
            cursor = conn.cursor()
            min_order = self.get_min_sort_order()
            
            # 创建 session，直接将 system 消息设为初始叶子节点
            cursor.execute(
                "INSERT INTO sessions (id, title, created_at, updated_at, has_messages, sort_order, current_leaf_msg_id) VALUES (?, ?, ?, ?, ?, ?, ?)",
                (session_id, title, now, now, 0, min_order, sys_msg_id)
            )
            
            cursor.execute("""
                INSERT INTO messages (id, session_id, role, content, reasoning, is_ignored, created_at, parent_id)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?)
            """, (sys_msg_id, session_id, "system", "你是一个优秀的助手！", "", 0, now, None))
            conn.commit()
            
            return dict(cursor.execute("SELECT * FROM sessions WHERE id = ?", (session_id,)).fetchone())

    # ------------------------------------------------------------------
    # 🆕 会话复制：深度克隆（全部分支 + 压缩标记 + 附件文件）
    # ------------------------------------------------------------------
    def _make_copy_title(self, base: str) -> str:
        """生成不重名的副本标题：X → X (副本) → X (副本 2) → …"""
        root = re.sub(r"\s*\(副本(?:\s*\d+)?\)\s*$", "", base or "").strip() or "新对话"
        with self.get_connection() as conn:
            existing = {r[0] for r in conn.execute("SELECT title FROM sessions").fetchall()}
        cand = f"{root} (副本)"
        n = 2
        while cand in existing and n < 1000:
            cand = f"{root} (副本 {n})"
            n += 1
        return cand

    def _dup_media_file(self, old_path: str, prefix: str):
        """把一个项目内媒体文件复制成新名字。
        返回 (新绝对路径, 新项目相对路径)；文件不存在或异常返回 None。"""
        if not old_path:
            return None
        old_abs = old_path if os.path.isabs(old_path) else os.path.join(self.files_root, old_path)
        if not os.path.isfile(old_abs):
            return None
        ext = os.path.splitext(old_abs)[1] or ".bin"
        new_dir = os.path.join(self.files_root, "data", "attachments")
        os.makedirs(new_dir, exist_ok=True)
        new_abs = ""
        for _ in range(5):
            new_abs = os.path.join(new_dir, f"{prefix}_{uuid.uuid4().hex[:8]}{ext}")
            if not os.path.exists(new_abs):
                break
        shutil.copy2(old_abs, new_abs)
        new_rel = os.path.relpath(new_abs, self.files_root).replace("\\", "/")
        return new_abs, new_rel

    def _copy_attachment_files(self, meta_json: Optional[str]) -> Optional[str]:
        """附件深拷贝：image/pdf 的磁盘文件复制改名并改写 local_path；
        text 类型正文内联在 JSON 里，无需处理。
        任何异常/文件缺失都原样返回，绝不让复制整体失败。"""
        if not meta_json:
            return meta_json
        try:
            meta = json.loads(meta_json)
        except Exception:
            return meta_json
        if not isinstance(meta, dict) or not meta.get("attachments"):
            return meta_json
        changed = False
        for att in meta.get("attachments") or []:
            if not isinstance(att, dict):
                continue
            try:
                if att.get("type") in ("image", "pdf") and att.get("local_path"):
                    got = self._dup_media_file(
                        att["local_path"], "img" if att["type"] == "image" else "pdf")
                    if got:
                        att["local_path"] = got[1]
                        if "abs_path" in att:
                            att["abs_path"] = got[0]
                        changed = True
                if att.get("type") == "pdf":
                    for im in att.get("images") or []:
                        if not isinstance(im, dict):
                            continue
                        got = self._dup_media_file(
                            im.get("abs_path") or im.get("local_path"), "pdfimg")
                        if got:
                            im["abs_path"] = got[0]
                            im["local_path"] = got[1]
                            changed = True
            except Exception as e:
                print(f"[DB] ⚠️ 附件深拷贝失败（保留原路径）: {e}", flush=True)
        return json.dumps(meta, ensure_ascii=False) if changed else meta_json

    def copy_session(self, session_id: str, new_title: Optional[str] = None,
                     copy_attachments: bool = True) -> Optional[Dict]:
        """📋 深度复制一个会话。

        - messages 全部重新生成 ID，parent_id / current_leaf_msg_id 全量重映射
          → 分支、压缩标记（role='compaction'）都原样保留
        - image/pdf 附件文件物理复制成新文件 → 副本自包含，删任意一方不影响另一方
        - 单事务写入；源会话零改动
        返回新会话 dict；源不存在返回 None。
        """
        now = int(time.time())
        n_att = 0
        with self.get_connection() as conn:
            cursor = conn.cursor()
            src_row = cursor.execute(
                "SELECT * FROM sessions WHERE id = ?", (session_id,)).fetchone()
            if not src_row:
                return None
            src = dict(src_row)
            rows = [dict(r) for r in cursor.execute(
                "SELECT * FROM messages WHERE session_id = ? ORDER BY created_at ASC, rowid ASC",
                (session_id,)).fetchall()]

            # ① 新 ID 映射（保持 comp_ / msg_ 前缀约定）
            idmap = {}
            for m in rows:
                pre = "comp_" if m.get("role") == "compaction" else "msg_"
                idmap[m["id"]] = f"{pre}{uuid.uuid4().hex[:16]}"

            # ② 标题（重名自动递增）
            title = new_title or self._make_copy_title(src.get("title") or "新对话")

            # ③ 消息 created_at 保相对间隔（同层兄弟排序不变）
            t0 = min((m.get("created_at") or 0) for m in rows) if rows else now

            new_sid = f"sess_{uuid.uuid4().hex[:12]}"
            new_leaf = idmap.get(src.get("current_leaf_msg_id"))
            if rows and not new_leaf:
                # 兜底：源叶子不在链上（数据异常）→ 取副本里时间最新的一条
                last = max(rows, key=lambda m: (m.get("created_at") or 0))
                new_leaf = idmap.get(last["id"])
                print(f"[DB] ⚠️ copy_session 源叶子异常，回退 leaf={new_leaf}", flush=True)

            # ④ 新会话（置列表顶部、不带星标、模式跟随源）
            cursor.execute(
                "INSERT INTO sessions (id, title, created_at, updated_at, has_messages, "
                "sort_order, is_starred, current_leaf_msg_id, mode) VALUES (?,?,?,?,?,?,?,?,?)",
                (new_sid, title, now, now, src.get("has_messages") or 0,
                 self.get_min_sort_order(), 0, new_leaf, src.get("mode")))

            # ⑤ 逐条复制消息（parent 重映射 + 附件深拷贝）
            for m in rows:
                meta = m.get("attachment_metadata")
                if copy_attachments and meta:
                    new_meta = self._copy_attachment_files(meta)
                    if new_meta != meta:
                        n_att += 1
                    meta = new_meta
                cursor.execute(
                    "INSERT INTO messages (id, session_id, role, content, reasoning, is_ignored, "
                    "created_at, attachment_metadata, parent_id, timeline, usage, stop_reason, error_message) "
                    "VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?)",
                    (idmap[m["id"]], new_sid, m.get("role"), m.get("content"),
                     m.get("reasoning"), m.get("is_ignored"),
                     now + ((m.get("created_at") or 0) - t0),
                     meta, idmap.get(m.get("parent_id")),
                     m.get("timeline"), m.get("usage"),
                     m.get("stop_reason"), m.get("error_message")))

            conn.commit()
            out = dict(cursor.execute(
                "SELECT * FROM sessions WHERE id = ?", (new_sid,)).fetchone())

        try:
            print(f"[DB] copy_session {session_id[:8]} → {new_sid[:8]} "
                  f"消息={len(rows)} 附件深拷贝={n_att} 标题={title}", flush=True)
        except Exception:
            pass
        return out

    def update_session_title(self, session_id: str, new_title: str):
        with self.get_connection() as conn:
            conn.execute("UPDATE sessions SET title = ?, updated_at = ? WHERE id = ?", (new_title, int(time.time()), session_id))
            conn.commit()

    def delete_session(self, session_id: str):
        with self.get_connection() as conn:
            conn.execute("DELETE FROM sessions WHERE id = ?", (session_id,))
            conn.commit()
            
    # 下方其它基本Session功能保持不变...
    def mark_session_has_messages(self, session_id: str):
        with self.get_connection() as conn:
            conn.execute("UPDATE sessions SET has_messages = 1 WHERE id = ?", (session_id,))
            conn.commit()
            
    def check_session_needs_title(self, session_id: str) -> bool:
        with self.get_connection() as conn:
            row = conn.execute("SELECT title, has_messages FROM sessions WHERE id = ?", (session_id,)).fetchone()
            if row: return row[0] == "新对话" and row[1] == 1
            return False

    def get_min_sort_order(self) -> int:
        with self.get_connection() as conn:
            row = conn.execute("SELECT MIN(sort_order) FROM sessions").fetchone()
            return (row[0] or 0) - 1
            
    def update_session_order(self, ordered_ids: list):
        with self.get_connection() as conn:
            for idx, sid in enumerate(ordered_ids):
                conn.execute("UPDATE sessions SET sort_order = ? WHERE id = ?", (idx, sid))
            conn.commit()
            
    def update_session_star(self, session_id: str, is_starred: bool):
        with self.get_connection() as conn:
            conn.execute("UPDATE sessions SET is_starred = ?, updated_at = ? WHERE id = ?", (1 if is_starred else 0, int(time.time()), session_id))
            conn.commit()
            
    def is_session_starred(self, session_id: str) -> bool:
        with self.get_connection() as conn:
            row = conn.execute("SELECT is_starred FROM sessions WHERE id = ?", (session_id,)).fetchone()
            return bool(row and row[0])


    # ==================== 🌟 核心：消息链表树操作 ====================

    def get_session_leaf(self, session_id: str) -> Optional[str]:
        """获取会话当前的叶子节点ID"""
        with self.get_connection() as conn:
            row = conn.execute("SELECT current_leaf_msg_id FROM sessions WHERE id = ?", (session_id,)).fetchone()
            return row[0] if row else None

    def update_session_leaf(self, session_id: str, leaf_msg_id: str):
        """切换时间线：手动更新当前会话的叶子节点"""
        with self.get_connection() as conn:
            conn.execute("UPDATE sessions SET current_leaf_msg_id = ?, updated_at = ? WHERE id = ?", 
                         (leaf_msg_id, int(time.time()), session_id))
            conn.commit()

    def get_message_chain(self, session_id: str) -> List[Dict]:
        """🚀 极客级递归拉取：顺藤摸瓜，只返回当前激活时间线上的消息！彻底断绝下游污染！"""
        leaf_id = self.get_session_leaf(session_id)
        if not leaf_id:
            return []

        chain = []
        with self.get_connection() as conn:
            cursor = conn.cursor()
            current_id = leaf_id
            seen = set()   # 🐛 环检测：数据异常（如自环标记）时快速退出，防主线程死循环卡死
            
            # 使用 Python 循环向上回溯（比 SQL CTE 更好调试，性能对于本地几千条聊天来说在 1ms 内）
            while current_id and current_id not in seen:
                seen.add(current_id)
                cursor.execute("SELECT * FROM messages WHERE id = ?", (current_id,))
                msg = cursor.fetchone()
                if not msg:
                    break
                chain.append(dict(msg))
                current_id = msg['parent_id']
            if current_id in seen:
                try:
                    print(f"[DB] ⚠️ get_message_chain 检测到环（session={session_id}），已截断", flush=True)
                except Exception:
                    pass
                
        # 因为是向上回溯，拉出来的链条是反的，最后翻转一下恢复正序
        chain.reverse()
        return chain

    def get_branch_info(self, parent_id: str) -> List[Dict]:
        """获取某一父节点下的所有子分支消息 (第二阶段用于UI渲染 '1/3')"""
        if not parent_id: return []
        with self.get_connection() as conn:
            return [dict(row) for row in conn.execute(
                "SELECT * FROM messages WHERE parent_id = ? ORDER BY created_at ASC", (parent_id,)
            ).fetchall()]

    def add_message(self, session_id: str, role: str, content: str, parent_id: str,
                    reasoning: str = "", is_ignored: bool = False,
                    msg_id: Optional[str] = None, attachment_metadata: Optional[str] = None,
                    timeline: Optional[str] = None,
                    usage: Optional[str] = None,
                    stop_reason: Optional[str] = None,
                    error_message: Optional[str] = None) -> Dict:
        """添加新消息，并自动将该消息设为当前会话的最新叶子节点
        🆕 P1: usage —— assistant 回复的精确 usage JSON（如 '{"input":..,"output":..}'）
        🆕 失败轮次: stop_reason/error_message —— 'error' 行入库但不回退叶子
                     （对照 pi：错误也持久化，回放时由 build_api_context 决定取舍）"""
        if not msg_id: msg_id = f"msg_{uuid.uuid4().hex[:16]}"
        now = int(time.time())
        
        with self.get_connection() as conn:
            cursor = conn.cursor()
            # 1. 插入消息
            cursor.execute("""
                INSERT INTO messages (id, session_id, role, content, reasoning, is_ignored, created_at, attachment_metadata, parent_id, timeline, usage, stop_reason, error_message)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """, (msg_id, session_id, role, content, reasoning, 1 if is_ignored else 0, now, attachment_metadata, parent_id, timeline, usage, stop_reason, error_message))
            
            # 2. 自动更新 session 的叶子节点（时间线前推）
            cursor.execute("UPDATE sessions SET current_leaf_msg_id = ?, updated_at = ? WHERE id = ?", 
                           (msg_id, now, session_id))
            conn.commit()
            try:
                print(f"[DB] add_message role={role} 内容={len(content or '')}c "
                      f"思考={len(reasoning or '')}c 时间线={'有' if timeline else '无'} "
                      f"id={msg_id} session={session_id[:8]}", flush=True)
            except Exception:
                pass
            
            return dict(cursor.execute("SELECT * FROM messages WHERE id = ?", (msg_id,)).fetchone())

    # ------------------------------------------------------------------
    # 🆕 压缩持久化：链上标记点（子叶子回溯到标记即停）
    # 插入后链形：…→ cut_before → [MARK role="compaction"] → first_retained → …→叶子
    # 历史行全部保留（前端切会话渲染照常）；API 上下文构建在标记处截断。
    # 关键：不动 sessions.current_leaf_msg_id（绝不能用 add_message）；
    #       INSERT + UPDATE 同一事务，不留断链窗口。
    # ------------------------------------------------------------------
    def insert_compaction_mark(self, session_id: str, summary: str,
                               cut_before_id: str, first_retained_id: str,
                               meta_json: Optional[str] = None) -> Optional[str]:
        """在链上插入压缩标记行并把保留首条的 parent_id 改指到标记。返回 mark_id。"""
        if not cut_before_id or not first_retained_id:
            return None
        if cut_before_id == first_retained_id:
            # 🐛 防自环：同一行不能既做切点前又做保留首条（timeline 回放同 id 场景）
            try:
                print(f"[DB] ⚠️ insert_compaction_mark 拒绝自环 cut==retained={cut_before_id}", flush=True)
            except Exception:
                pass
            return None
        mark_id = f"comp_{uuid.uuid4().hex[:16]}"
        now = int(time.time())
        with self.get_connection() as conn:
            cursor = conn.cursor()
            cursor.execute("""INSERT INTO messages
                (id, session_id, role, content, reasoning, is_ignored,
                 created_at, attachment_metadata, parent_id, timeline, usage)
                VALUES (?, ?, 'compaction', ?, '', 1, ?, ?, ?, NULL, NULL)""",
                           (mark_id, session_id, summary, now, meta_json, cut_before_id))
            cursor.execute("UPDATE messages SET parent_id = ? WHERE id = ?",
                           (mark_id, first_retained_id))
            conn.commit()
            try:
                print(f"[DB] insert_compaction_mark session={session_id[:8]} "
                      f"cut_before={cut_before_id} first_retained={first_retained_id} "
                      f"mark={mark_id} summary={len(summary or '')}c", flush=True)
            except Exception:
                pass
        return mark_id
        # === 在 db_manager.py 中添加这个方法 ===
    def get_branch_leaf(self, msg_id: str) -> str:
        """寻找一条时间线的最末端叶子节点"""
        current_id = msg_id
        with self.get_connection() as conn:
            cursor = conn.cursor()
            while True:
                # 寻找把当前节点作为父节点的子节点，按时间倒序取最新的一条（意味着它顺着最近被聊过的那条线往下走）
                cursor.execute("SELECT id FROM messages WHERE parent_id = ? ORDER BY created_at DESC LIMIT 1", (current_id,))
                child = cursor.fetchone()
                if child:
                    current_id = child[0]
                else:
                    break  # 没有子节点了，它自己就是叶子！
        return current_id
    def delete_message_branch(self, session_id: str, msg_id: str):
        """🚀 精准剪枝：删AI只删当前分支，删User连根拔起，并自动平滑回退时间线"""
        with self.get_connection() as conn:
            cursor = conn.cursor()
            
            # 1. 查出要删除的节点的父亲
            cursor.execute("SELECT parent_id FROM messages WHERE id = ?", (msg_id,))
            row = cursor.fetchone()
            if not row: return
            safe_parent_id = row['parent_id']

            # 2. 目标就是传进来的 msg_id 本身 (不再强制上移到 parent)
            target_id = msg_id

            # 3. 拉取全会话，构建亲属关系树
            cursor.execute("SELECT id, parent_id FROM messages WHERE session_id = ?", (session_id,))
            all_msgs = cursor.fetchall()
            
            children_map = {}
            for m in all_msgs:
                pid = m['parent_id']
                if pid not in children_map:
                    children_map[pid] = []
                children_map[pid].append(m['id'])
            
            # 4. 广度优先搜索 (BFS)：找出目标消息及其所有子孙
            to_delete = set([target_id])
            queue = [target_id]
            while queue:
                curr = queue.pop(0)
                if curr in children_map:
                    for child in children_map[curr]:
                        if child not in to_delete:
                            to_delete.add(child)
                            queue.append(child)
            
            # 5. 判断当前时间线的“叶子节点”是否在被波及的名单里
            cursor.execute("SELECT current_leaf_msg_id FROM sessions WHERE id = ?", (session_id,))
            leaf_row = cursor.fetchone()
            leaf_needs_update = leaf_row and leaf_row['current_leaf_msg_id'] in to_delete

            # 6. 🌟 执行物理删除前，收集将被删除的附件元数据
            deleted_metadata = []
            for d_id in to_delete:
                # 先查出它的 metadata
                cursor.execute("SELECT attachment_metadata FROM messages WHERE id = ?", (d_id,))
                row = cursor.fetchone()
                if row and row['attachment_metadata']:
                    deleted_metadata.append(row['attachment_metadata'])
                
                # 然后再执行物理删除
                cursor.execute("DELETE FROM messages WHERE id = ?", (d_id,))

            
            # 7. 🌟 核心：如果时间线断了，自动寻找平滑降落点
            sibling_row = None
            if leaf_needs_update:
                # 尝试寻找被删节点的最新“兄弟姐妹” (例如删了分支2，寻找分支1)
                cursor.execute("SELECT id FROM messages WHERE parent_id = ? ORDER BY created_at DESC LIMIT 1", (safe_parent_id,))
                sibling_row = cursor.fetchone()
                
                # 如果有兄弟，降落到兄弟；如果没兄弟(只有1次回答)，退回原点(提问)
                new_leaf = sibling_row['id'] if sibling_row else safe_parent_id
                cursor.execute("UPDATE sessions SET current_leaf_msg_id = ? WHERE id = ?", (new_leaf, session_id))
            
            conn.commit()

        # 8. 如果降落到了兄弟分支，兄弟可能还有下文，需再次对齐真实叶子节点
        if leaf_needs_update and sibling_row:
            real_leaf = self.get_branch_leaf(new_leaf)
            self.update_session_leaf(session_id, real_leaf)
            
        return deleted_metadata  # 🌟 返回被删除的元数据，交给 MainWindow 去粉碎文件