# -*- coding: utf-8 -*-
"""
wxwork_db.py - 企业微信本地数据库读取模块（引擎 B 升级数据源）
=================================================================
基于 wxwork_crypto 解密 message.db / session.db / user.db 为标准 SQLite，
提供增量消息轮询接口，供 engine_b 使用。

参考实现（C:\\Users\\pc\\WorkBuddy\\2026-08-19-18-27-34\\，密钥已验证正确）：
  - wxwork_keys.json          已复制的验证过密钥（本项目 wxwork_keys.json）
  - wxwork_export_final.py    移植：parse_content / extract_protobuf_texts /
                              clean_text / MSG_TYPE_MAP / is_personal_chat /
                              is_blocked_conv_name / format_timestamp /
                              load_keys / detect_wxwork_dir / connect_sqlite /
                              decrypt_with_keys（WAL 感知增量解密缓存）/
                              会话名解析（user_cache / conv_cache / S: 对端推导）
  - wxwork_crypto.py          解密算法与本项目 wxwork_crypto.py 完全一致（互相印证）

相对参考实现的增强：
  - WAL 帧合并：企微运行时新消息在 message.db-wal（未 checkpoint），参考实现仅
    以 wal_mtime 触发重解密主库（仍缺 WAL 内最新消息）。本项目直接解密 WAL 帧
    （帧数据与主库同算法加密，帧头明文），按页号覆盖合并进主库解密副本，
    经 PRAGMA quick_check 验证（见 test_wal_merge.py）。
  - 会话指纹：复用 wechat_bot.WeChatBot._fp_from_name 的名称哈希算法，
    使 DB 直读的 fp_hex 与引擎 A（identity_by_name 模式）完全一致，
    引擎 A 可在会话列表按名称重新定位窗口。

典型流程:
  keys = load_keys()
  db = WXWorkDB(detect_wxwork_dir(), keys)
  msgs = db.get_new_messages(since_ts=...)
  db.close()
"""
from __future__ import annotations

import hashlib
import json
import os
import re
import sqlite3
import struct
import time
import unicodedata
from datetime import datetime

from wxwork_crypto import (
    decrypt_db_to_file,
    decrypt_page,
    read_page_header,
    SQLITE_FILE_HEADER,
    verify_key,
)

_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
KEYS_FILE = os.path.join(_SCRIPT_DIR, "wxwork_keys.json")
DEFAULT_CACHE_DIR = os.path.join(_SCRIPT_DIR, "wxwork_decrypted")

# 企微数据库文件名（与聊天相关的库）
_RELEVANT_DBS = (
    "message.db", "session.db", "user.db", "company.db",
    "message_lookup.db", "user_extend.db",
)

# 引擎 A 会话指纹字节数（wechat_bot.py: _AVATAR_FP_BYTES=8 + _NAME_FP_BYTES=32）
SESSION_FP_BYTES = 40

# ---------------------------------------------------------------------------
# 消息类型映射 / 会话过滤（移植自 wxwork_export_final.py）
# ---------------------------------------------------------------------------
MSG_TYPE_MAP = {
    0: "文本", 1: "文本", 2: "文本",
    3: "图片", 4: "语音", 5: "表情",
    6: "链接", 7: "文件", 8: "视频",
    9: "位置", 10: "名片", 11: "系统",
    12: "引用", 13: "红包", 14: "图片",
    15: "转账", 16: "语音", 17: "视频号",
    20: "合并转发", 21: "日程", 22: "接龙",
    23: "文件回复", 26: "位置共享", 29: "视频通话",
    31: "图文链接", 38: "系统通知", 40: "待办",
    42: "文件预览", 46: "话题", 47: "图文",
    49: "待办", 51: "收藏", 53: "订阅通知",
    55: "文件分享", 56: "群公告", 57: "投票",
    59: "回执", 61: "文件编辑", 62: "文件评论",
    63: "收藏合并", 64: "卡片", 65: "群文件",
    66: "工作台", 67: "位置共享", 68: "视频通话",
    69: "音频通话", 70: "企业微信应用", 73: "微信好友",
    74: "回复", 76: "邀请", 77: "移除",
    78: "修改群名", 79: "修改群公告", 80: "加入群聊",
    81: "退出群聊", 82: "解散群聊", 83: "群主转让",
    101: "系统", 111: "文件", 123: "截图",
    132: "系统",
    503: "应用消息", 529: "朋友圈", 561: "应用消息",
    565: "应用消息", 573: "应用消息",
    1002: "安全通知", 1011: "系统消息", 1012: "系统消息",
    1017: "系统消息", 1022: "系统消息", 1025: "系统消息",
    1043: "系统消息",
    1988: "系统消息",
}


def get_msg_type_name(ct):
    """消息类型名"""
    try:
        ct = int(ct)
    except (ValueError, TypeError):
        return f"类型{ct}"
    return MSG_TYPE_MAP.get(ct, f"类型{ct}")


def is_personal_chat(conv_id):
    """判断会话是否为人与人之间的单聊（移植自 wxwork_export_final.py）。

    保留: M:xxx（微信单聊）、S:UID_UID（企微单聊）
    过滤: R:xxx 群聊 / Y:xxx 应用 / O:xxx 第三方 / MAIL/APPROVAL 系统虚拟会话
    """
    c = str(conv_id or "")
    if c.startswith("M:"):
        return True
    if c.startswith("S:"):
        parts = c[2:].split("_")
        if len(parts) == 2 and parts[0].isdigit() and parts[1].isdigit():
            return True
        return False
    return False


# 按会话名称过滤的官方/系统账号
DEFAULT_BLOCKED_CONV_NAMES = ("企业微信团队", "微信团队", "微信支付", "腾讯客服", "腾讯新闻")


def is_blocked_conv_name(conv_name, blocked=None):
    """按会话名称判断是否命中过滤名单（官方/系统账号）。"""
    if not conv_name:
        return False
    blocked = DEFAULT_BLOCKED_CONV_NAMES if blocked is None else tuple(blocked)
    return any(kw and kw in str(conv_name) for kw in blocked)


# ---------------------------------------------------------------------------
# 消息内容解析（移植自 wxwork_export_final.py）
# ---------------------------------------------------------------------------
def clean_text(text):
    """清洗提取的文本: 去除控制字符和 protobuf 残留。"""
    if not text:
        return ""
    cleaned = "".join(c for c in text if c.isprintable() or c in "\n\t")
    cleaned = cleaned.strip()
    cleaned = re.sub(r"[ \t]+", " ", cleaned)
    cleaned = re.sub(r"\n{3,}", "\n\n", cleaned)
    return cleaned


def extract_protobuf_texts(data):
    """从 protobuf 编码的 content 中递归提取可读文本，返回去重后的文本列表。"""
    if not data or not isinstance(data, bytes):
        return []

    texts = []
    seen = set()

    def add_text(t):
        t = clean_text(t)
        if len(t) >= 2 and t not in seen:
            has_chinese = any("\u4e00" <= c <= "\u9fff" for c in t)
            alpha_count = sum(1 for c in t if c.isalpha())
            if has_chinese or alpha_count >= 6:
                seen.add(t)
                texts.append(t)

    def parse(buf, depth):
        if depth > 6 or not buf:
            return
        i = 0
        while i < len(buf):
            try:
                tag = 0
                shift = 0
                while i < len(buf) and shift < 70:
                    b = buf[i]
                    tag |= (b & 0x7f) << shift
                    i += 1
                    if not (b & 0x80):
                        break
                    shift += 7
                if shift >= 70:
                    break
                field_num = tag >> 3
                wire_type = tag & 7
                if field_num == 0:
                    break

                if wire_type == 0:  # varint
                    while i < len(buf):
                        b = buf[i]
                        i += 1
                        if not (b & 0x80):
                            break
                elif wire_type == 1:  # 64-bit
                    i += 8
                elif wire_type == 2:  # length-delimited
                    length = 0
                    shift = 0
                    while i < len(buf) and shift < 70:
                        b = buf[i]
                        length |= (b & 0x7f) << shift
                        i += 1
                        if not (b & 0x80):
                            break
                        shift += 7
                    if shift >= 70:
                        break
                    chunk = buf[i:i + length]
                    i += length
                    try:
                        text = chunk.decode("utf-8")
                        if all(ord(c) >= 32 or c in "\n\t\r" for c in text):
                            add_text(text)
                    except (UnicodeDecodeError, ValueError):
                        pass
                    if length < 8192:
                        parse(chunk, depth + 1)
                elif wire_type == 5:  # 32-bit
                    i += 4
                else:
                    break
            except Exception:
                break

    parse(data, 0)

    if not any("\u4e00" <= c <= "\u9fff" for t in texts for c in t[:50]):
        try:
            text = data.decode("utf-8", errors="ignore")
            readable = re.findall(
                r"[\u4e00-\u9fff\w\s，。！？、；：“”‘’（）《》【】\-—.,;:!?()\[\]/\\+*=<>@#$%^&~`]{2,}",
                text)
            for r in readable:
                add_text(r.strip())
        except Exception:
            pass

    return texts


def parse_content(content):
    """解析消息内容, 返回可读文本（移植自 wxwork_export_final.py）。"""
    if content is None:
        return ""
    if not isinstance(content, bytes):
        content = str(content).encode("utf-8", errors="replace")
    if not content:
        return ""

    # 情况 1: 直接是 UTF-8 文本
    try:
        text = content.decode("utf-8")
        if text and all(ord(c) >= 32 or c in "\n\r\t" for c in text[:500]):
            return clean_text(text)
    except (UnicodeDecodeError, ValueError):
        pass

    # 情况 2: protobuf 编码
    texts = extract_protobuf_texts(content)
    if texts:
        texts.sort(key=len, reverse=True)
        meaningful = [t for t in texts
                      if not (t.startswith("{") and "}" in t)
                      and not t.startswith("http")
                      and len(t) > 1]
        if meaningful:
            cleaned_msgs = []
            for t in meaningful:
                t = t.strip()
                m = re.match(r"^(.{6,}?)[?zxa-zA-Z]{1,4}\1$", t)
                if m:
                    t = m.group(1)
                t = re.sub(
                    r"^[?zxa-zA-Z]{1,4}(?=[\u4e00-\u9fff\u3000-\u303f\uff00-\uffef])", "", t)
                t = re.sub(
                    r"^[0-9](?=[\u4e00-\u9fff])(?![\u4e00-\u9fff]{0,2}"
                    r"(?:月|日|年|个|天|点|号|楼|人|次|分|秒|周|元|块|台|家|辆|条|件|张|本|层|岁|万|亿|美元|室|房|折))",
                    "", t)
                t = t.strip()
                if len(t) >= 3:
                    dup = False
                    for c in cleaned_msgs:
                        if len(c) >= 4 and (t in c or c in t):
                            dup = True
                            break
                    if dup:
                        continue
                if t and t not in cleaned_msgs:
                    cleaned_msgs.append(t)
            if cleaned_msgs:
                return " | ".join(cleaned_msgs[:2])
        return texts[0]

    # 情况 3: JSON
    try:
        obj = json.loads(content.decode("utf-8", errors="replace"))
        return json.dumps(obj, ensure_ascii=False)[:2000]
    except Exception:
        pass

    # 情况 4: 二进制数据
    if len(content) <= 64:
        return content.hex()
    return f"[二进制数据 {len(content)} 字节]"


def format_timestamp(ts):
    """Unix 时间戳转可读格式（>10^12 视为毫秒）。"""
    if ts is None or ts == "":
        return ""
    try:
        ts = int(ts)
        if ts > 10 ** 12:
            ts = ts // 1000
        if ts < 0:
            return str(ts)
        return datetime.fromtimestamp(ts).strftime("%Y-%m-%d %H:%M:%S")
    except (ValueError, OSError, OverflowError):
        return str(ts)


# ---------------------------------------------------------------------------
# 会话指纹（与 wechat_bot.WeChatBot._fp_from_name 算法一致）
# ---------------------------------------------------------------------------
_WHITESPACE_RE = re.compile(r"\s+")


def normalize_name(value: str) -> str:
    """昵称归一化（与 session_name.normalize_name 完全一致）。

    实测企微标题读出来是「高瑞 @微信」、列表行是「高瑞@微信」，中间多个空格；
    不归一的话同一个人会算出两个 md5。仅用标准库实现（session_name 顶层
    依赖 numpy，DB 模块不引入视觉依赖）。
    """
    text = unicodedata.normalize("NFKC", str(value or ""))
    text = _WHITESPACE_RE.sub("", text)
    return text.strip()


def session_fp_from_name(name: str) -> str:
    """会话名 → 40 字节会话指纹 hex。

    与 wechat_bot.py::_fp_from_name 完全一致：前 16 字节 = md5(归一化昵称)，
    后 24 字节 = blake2b(md5)。引擎 A 在 identity_by_name 模式下用同一算法
    从屏幕 OCR 的昵称生成指纹，因此 DB 直读投递的 fp_hex 可被引擎 A
    在会话列表中按名称重新定位窗口。
    """
    normalized = normalize_name(name)
    if not normalized:
        return ""
    digest = hashlib.md5(normalized.encode("utf-8")).digest()
    tail = hashlib.blake2b(digest, digest_size=SESSION_FP_BYTES - len(digest)).digest()
    return (digest + tail).hex()


# ---------------------------------------------------------------------------
# 密钥与解密（移植自 wxwork_export_final.py + WAL 帧合并增强）
# ---------------------------------------------------------------------------
def load_keys() -> dict:
    """加载密钥映射 {user_dir: key_hex}，支持 global_key/* 通配。"""
    if not os.path.exists(KEYS_FILE):
        return {}
    with open(KEYS_FILE, encoding="utf-8") as f:
        data = json.load(f)
    keys = data.get("keys", {})
    if "global_key" in data:
        keys.setdefault("*", data["global_key"])
    return keys


def detect_wxwork_dir(candidates=None) -> str | None:
    """自动检测企业微信数据根目录（移植自 wxwork_export_final.py）。

    依次尝试: 配置候选 → 当前用户 Documents\\WXWork →
    注册表 Personal 下的 WXWork → 扫描所有用户。
    返回第一个含账号 Data/message.db 的有效目录。
    """
    def _valid(base):
        if not base or not os.path.isdir(base):
            return False
        try:
            for d in os.listdir(base):
                if os.path.exists(os.path.join(base, d, "Data", "message.db")):
                    return True
        except OSError:
            pass
        return False

    for c in (candidates or []):
        if _valid(c):
            return c
    default = os.path.join(os.path.expanduser("~"), "Documents", "WXWork")
    if _valid(default):
        return default
    try:
        import winreg
        with winreg.OpenKey(
            winreg.HKEY_CURRENT_USER,
            r"Software\Microsoft\Windows\CurrentVersion\Explorer\Shell Folders",
        ) as k:
            personal, _ = winreg.QueryValueEx(k, "Personal")
        if personal:
            reg_path = os.path.join(personal, "WXWork")
            if _valid(reg_path):
                return reg_path
    except Exception:
        pass
    users_root = os.path.join(os.environ.get("SystemDrive", "C:"), os.sep, "Users")
    if os.path.isdir(users_root):
        try:
            for u in sorted(os.listdir(users_root)):
                p = os.path.join(users_root, u, "Documents", "WXWork")
                if _valid(p):
                    return p
        except OSError:
            pass
    return None


def _is_plain_sqlite(path: str) -> bool:
    """文件头是否为明文 SQLite。"""
    try:
        with open(path, "rb") as f:
            return f.read(16) == SQLITE_FILE_HEADER
    except OSError:
        return False


def wal_read_frames(wal_path: str, key: bytes) -> tuple[dict, dict]:
    """读取 WAL 帧并解密，返回 ({pgno: bytes}, meta)。

    帧数据与主库同算法加密（wxSQLite3 aes128cbc，帧头 24 字节明文）。
    仅应用 salt 与 WAL 头一致的帧（跳过已回滚/旧事务帧），后帧覆盖前帧。
    """
    pages: dict[int, bytes] = {}
    meta = {"pgsz": 0, "frames": 0, "applied": 0, "salt_ok": 0}
    try:
        with open(wal_path, "rb") as f:
            hdr = f.read(32)
            if len(hdr) < 32:
                return pages, meta
            magic, ver, pgsz, ckpt, salt1, salt2, _, _ = struct.unpack(">IIIIIIII", hdr)
            meta["pgsz"] = pgsz
            while True:
                fh = f.read(24)
                if len(fh) < 24:
                    break
                pgno, commit, fs1, fs2, _, _ = struct.unpack(">IIIIII", fh)
                data = f.read(pgsz)
                if len(data) < pgsz:
                    break
                meta["frames"] += 1
                if fs1 != salt1 or fs2 != salt2:
                    continue
                meta["salt_ok"] += 1
                if pgno == 0 or pgno > (1 << 31):
                    continue
                dec = decrypt_page(key, pgno, data, pgsz)
                if dec is None:
                    continue
                pages[pgno] = dec
                meta["applied"] += 1
    except OSError:
        pass
    return pages, meta


def merge_wal_into_db(db_path: str, wal_path: str, key: bytes, out_path: str) -> bool:
    """解密主库 + 合并 WAL 帧 → out_path。返回是否成功。

    主库与 WAL 帧都解密后，按页号覆盖（等价于 checkpoint 后的状态），
    并更新 SQLite 头的页数字段。WAL 无有效帧时退化为仅主库。
    """
    if not decrypt_db_to_file(db_path, key, out_path):
        return False
    pages, meta = wal_read_frames(wal_path, key)
    if not pages:
        return True
    pgsz = meta["pgsz"] or 4096
    try:
        with open(out_path, "r+b") as f:
            size = os.path.getsize(out_path)
            n_pages = (size + pgsz - 1) // pgsz
            for pgno, data in pages.items():
                if pgno > n_pages:
                    f.seek(0, os.SEEK_END)
                    f.write(b"\x00" * ((pgno - n_pages) * pgsz))
                    n_pages = pgno
                f.seek((pgno - 1) * pgsz)
                f.write(data)
            if n_pages > 0:
                f.seek(28)
                f.write(struct.pack(">I", n_pages))
    except OSError:
        return False
    return True


def decrypt_with_keys(db_base: str, out_dir: str, keys_map: dict, use_cache: bool = True) -> list:
    """解密所有可解密的账号数据库（移植自 wxwork_export_final.py）。

    返回 [(out_path, db_name, user_dir), ...]。
    - 增量缓存: 解密副本 mtime >= max(db, wal) 且大小正常则复用
    - WAL 增强: message.db 的 WAL 非空时合并帧（拿 checkpoint 前的最新消息）
    - 密钥: 优先目录密钥，失败尝试全部密钥
    """
    decrypted_dbs = []
    os.makedirs(out_dir, exist_ok=True)

    all_keys = set()
    for k in keys_map.values():
        try:
            all_keys.add(bytes.fromhex(k.replace("x'", "").replace("'", "")))
        except ValueError:
            pass

    if not os.path.isdir(db_base):
        return decrypted_dbs

    for user_dir in sorted(os.listdir(db_base)):
        data_dir = os.path.join(db_base, user_dir, "Data")
        if not os.path.isdir(data_dir):
            continue
        db_files = [f for f in os.listdir(data_dir)
                    if f.endswith(".db") and not f.endswith("-wal") and not f.endswith("-shm")]
        if not db_files:
            continue

        user_out = os.path.join(out_dir, user_dir)
        os.makedirs(user_out, exist_ok=True)

        dir_key = None
        if user_dir in keys_map:
            try:
                dir_key = bytes.fromhex(keys_map[user_dir])
            except ValueError:
                dir_key = None

        for db_name in _RELEVANT_DBS:
            db_path = os.path.join(data_dir, db_name)
            if not os.path.isfile(db_path):
                continue
            out_path = os.path.join(user_out, db_name)

            # 增量缓存判断（WAL 感知）
            wal_path = db_path + "-wal"
            wal_mtime = 0
            try:
                if os.path.exists(wal_path) and os.path.getsize(wal_path) > 0:
                    wal_mtime = os.path.getmtime(wal_path)
            except OSError:
                wal_mtime = 0
            src_mtime = max(os.path.getmtime(db_path), wal_mtime)

            if use_cache and os.path.exists(out_path):
                try:
                    if (os.path.getmtime(out_path) >= src_mtime
                            and os.path.getsize(out_path) > 4096):
                        decrypted_dbs.append((out_path, db_name, user_dir))
                        continue
                except OSError:
                    pass

            # 明文库直接复制
            if _is_plain_sqlite(db_path):
                try:
                    with open(db_path, "rb") as fin, open(out_path, "wb") as fout:
                        fout.write(fin.read())
                    decrypted_dbs.append((out_path, db_name, user_dir))
                    continue
                except OSError:
                    continue

            # 加密库：目录密钥 → 全部密钥
            if read_page_header(db_path) is None:
                continue  # 不是企微加密格式
            candidate_keys = [dir_key] if dir_key else []
            candidate_keys += [k for k in all_keys if k != dir_key]
            decrypted_ok = False
            for key in candidate_keys:
                if key is None:
                    continue
                try:
                    if not verify_key(key, db_path):
                        continue
                    # message.db 且 WAL 非空 → 合并；否则整库解密
                    if db_name == "message.db" and wal_mtime > 0:
                        ok = merge_wal_into_db(db_path, wal_path, key, out_path)
                    else:
                        ok = decrypt_db_to_file(db_path, key, out_path)
                    if ok:
                        decrypted_dbs.append((out_path, db_name, user_dir))
                        decrypted_ok = True
                        break
                except Exception:
                    continue
            if not decrypted_ok:
                try:
                    os.remove(out_path)
                except OSError:
                    pass

    return decrypted_dbs


def connect_sqlite(path: str) -> sqlite3.Connection:
    conn = sqlite3.connect(f"file:{path}?mode=ro", uri=True)
    conn.text_factory = lambda b: b.decode("utf-8", errors="replace")
    return conn


# ---------------------------------------------------------------------------
# WXWorkDB：多账号增量消息读取器
# ---------------------------------------------------------------------------
class WXWorkDB:
    """企微数据根目录下的加密数据库读取器（只读，多账号）。

    - 构造时: 验证密钥 → 确保解密缓存（WAL 感知）→ 加载元数据 → 打开各账号
      message.db 解密副本连接
    - 轮询: get_new_messages(since_ts) 返回 since_ts 之后的新客户消息；
      内部按 mtime 检查 WAL/主库变化，必要时重解密并重开连接
    """

    def __init__(self, db_base: str, keys_map: dict | None = None,
                 cache_dir: str | None = None, poll_interval: float = 2.0):
        self.db_base = db_base
        self.keys_map = keys_map if keys_map is not None else load_keys()
        self.cache_dir = cache_dir or DEFAULT_CACHE_DIR
        self.poll_interval = max(0.5, float(poll_interval))
        self.last_error = ""
        # (user_dir) -> sqlite3.Connection（message.db 解密副本）
        self._conns: dict[str, sqlite3.Connection] = {}
        # (user_dir) -> 解密副本路径
        self._paths: dict[str, str] = {}
        # 元数据缓存
        self.user_cache: dict = {}
        self.conv_cache: dict = {}
        self._open()

    # ── 生命周期 ────────────────────────────────────────────────────────────
    def _open(self):
        decrypted = decrypt_with_keys(self.db_base, self.cache_dir, self.keys_map)
        self._decrypted = decrypted
        if not decrypted:
            self.last_error = "没有可用密钥解密的账号数据库"
            return
        self._load_metadata()
        for out_path, db_name, user_dir in decrypted:
            if db_name != "message.db":
                continue
            try:
                conn = connect_sqlite(out_path)
                self._conns[user_dir] = conn
                self._paths[user_dir] = out_path
            except sqlite3.Error as exc:
                self.last_error = f"打开 {user_dir}/message.db 失败: {exc}"

    def _load_metadata(self):
        """加载 user_cache / conv_cache（移植自 wxwork_export_final.py 阶段 1）。"""
        for out_path, db_name, user_dir in self._decrypted:
            try:
                conn = connect_sqlite(out_path)
                cur = conn.cursor()
                if db_name == "user.db":
                    try:
                        cur.execute("PRAGMA table_info(user_table)")
                        cols = [r[1] for r in cur.fetchall()]
                        if cols:
                            cur.execute(
                                "SELECT id, name, real_name, account FROM user_table")
                            for uid, name, real_name, account in cur.fetchall():
                                if uid is not None:
                                    nm = name or real_name or account or str(uid)
                                    self.user_cache[(user_dir, str(uid))] = str(nm)
                    except sqlite3.Error:
                        pass
                elif db_name == "session.db":
                    try:
                        cur.execute("PRAGMA table_info(conversation_table)")
                        cols = [r[1] for r in cur.fetchall()]
                        if cols:
                            cur.execute(
                                "SELECT id, name, roomname_remark, session_id "
                                "FROM conversation_table")
                            for cid, name, remark, sid in cur.fetchall():
                                if cid:
                                    nm = remark or name or sid
                                    self.conv_cache[(user_dir, str(cid))] = (
                                        str(nm) if nm else "")
                    except sqlite3.Error:
                        pass
                conn.close()
            except sqlite3.Error:
                continue

    def _conv_display_name(self, user_dir: str, conv_id) -> str:
        """会话显示名（移植自 wxwork_export_final.py 的 conv_name 解析）。"""
        conv_id = str(conv_id or "")
        name = self.conv_cache.get((user_dir, conv_id), "")
        if name:
            return name
        if conv_id.startswith("M:"):
            uid = conv_id[2:]
            return self.user_cache.get((user_dir, uid), uid)
        if conv_id.startswith("S:"):
            parts = conv_id[2:].split("_")
            peer_uid = ""
            for p in parts:
                if p != user_dir:
                    peer_uid = p
                    break
            if not peer_uid:
                peer_uid = parts[0] if parts else ""
            peer_name = self.user_cache.get((user_dir, peer_uid), "")
            return peer_name if peer_name else ("单聊 " + conv_id)
        if conv_id.startswith("O:"):
            return conv_id[2:]
        if conv_id.startswith("Y:"):
            return "应用 " + conv_id[2:]
        return conv_id

    def _refresh_if_needed(self):
        """WAL/主库比解密副本新 → 重解密并重开连接。"""
        for out_path, db_name, user_dir in self._decrypted:
            if db_name != "message.db":
                continue
            src = os.path.join(self.db_base, user_dir, "Data", db_name)
            wal = src + "-wal"
            try:
                src_mtime = os.path.getmtime(src)
                wal_mtime = 0
                if os.path.exists(wal) and os.path.getsize(wal) > 0:
                    wal_mtime = os.path.getmtime(wal)
                src_mtime = max(src_mtime, wal_mtime)
                if os.path.getmtime(out_path) >= src_mtime:
                    continue
            except OSError:
                continue
            # 重解密
            key = None
            for candidate in self._keys_for(user_dir):
                try:
                    if verify_key(candidate, src):
                        key = candidate
                        break
                except Exception:
                    continue
            if key is None:
                continue
            try:
                if wal_mtime > 0:
                    ok = merge_wal_into_db(src, wal, key, out_path)
                else:
                    ok = decrypt_db_to_file(src, key, out_path)
                if ok:
                    old = self._conns.pop(user_dir, None)
                    if old is not None:
                        old.close()
                    self._conns[user_dir] = connect_sqlite(out_path)
                    self._paths[user_dir] = out_path
            except Exception as exc:
                self.last_error = f"重解密 {user_dir} 失败: {exc}"

    def _keys_for(self, user_dir: str) -> list:
        keys = []
        if user_dir in self.keys_map:
            try:
                keys.append(bytes.fromhex(self.keys_map[user_dir]))
            except ValueError:
                pass
        if "*" in self.keys_map:
            try:
                keys.append(bytes.fromhex(self.keys_map["*"]))
            except ValueError:
                pass
        return keys

    # ── 消息读取 ────────────────────────────────────────────────────────────
    def get_new_messages(self, since_ts: float) -> list[dict]:
        """返回 since_ts（秒）之后的新客户消息列表。

        返回字段:
          account    账号目录名
          conv_id    企微会话 ID（M:/S:/R:...）
          fp_hex     会话名派生的 40 字节指纹 hex（引擎 A identity_by_name 兼容）
          display_name 会话显示名
          sender_id  发送者 UID
          is_self    是否本账号自己发送
          content    解析后的可读文本（空文本不返回）
          content_type 原始消息类型
          send_time  秒级时间戳
          server_id  服务器消息 ID（去重键）
          dedup_key  f"{fp_hex}:{send_time}"
        """
        self._refresh_if_needed()
        result: list[dict] = []
        try:
            since = float(since_ts or 0.0)
        except (TypeError, ValueError):
            since = 0.0
        for user_dir, conn in self._conns.items():
            try:
                cur = conn.cursor()
                cur.execute(
                    "PRAGMA table_info(message_table)")
                columns = [row[1] for row in cur.fetchall()]
                if "send_time" not in columns or "conversation_id" not in columns:
                    continue
                # 增量查询（send_time 为秒级 Unix 时间戳）
                cur.execute(
                    "SELECT * FROM message_table WHERE send_time > ? "
                    "ORDER BY send_time ASC LIMIT 2000",
                    (since,))
                for row in cur.fetchall():
                    msg = dict(zip(columns, row))
                    m = self._parse_message(user_dir, msg)
                    if m:
                        result.append(m)
            except sqlite3.Error as exc:
                self.last_error = f"查询 {user_dir} 失败: {exc}"
        result.sort(key=lambda m: m["send_time"])
        return result

    def _parse_message(self, user_dir: str, msg: dict) -> dict | None:
        conv_id = str(msg.get("conversation_id") or "")
        # 仅人与人单聊（过滤群聊/应用/第三方/系统虚拟会话）
        if not is_personal_chat(conv_id):
            return None
        sender_id = str(msg.get("sender_id") or "")
        send_time = msg.get("send_time") or 0
        try:
            send_time = float(send_time)
        except (TypeError, ValueError):
            return None
        if send_time <= 0:
            return None
        # 本账号自己发送的消息不投递（引擎 B 只管客户消息）
        is_self = bool(sender_id) and sender_id == str(user_dir)
        content = parse_content(msg.get("content"))
        if not content.strip():
            return None
        display_name = self._conv_display_name(user_dir, conv_id)
        if is_blocked_conv_name(display_name) or is_blocked_conv_name(conv_id):
            return None
        fp_hex = session_fp_from_name(display_name)
        if not fp_hex:
            return None
        server_id = str(msg.get("server_id") or "")
        return {
            "account": user_dir,
            "conv_id": conv_id,
            "fp_hex": fp_hex,
            "display_name": display_name,
            "sender_id": sender_id,
            "is_self": is_self,
            "content": content,
            "content_type": msg.get("content_type"),
            "send_time": send_time,
            "server_id": server_id,
            "dedup_key": f"{fp_hex}:{send_time}",
        }

    def health_check(self) -> bool:
        """连接与密钥有效性自检。"""
        if not self._conns:
            return False
        try:
            for conn in self._conns.values():
                cur = conn.cursor()
                cur.execute("SELECT COUNT(*) FROM sqlite_master")
                cur.fetchone()
            return True
        except Exception:
            return False

    def close(self):
        for conn in self._conns.values():
            try:
                conn.close()
            except Exception:
                pass
        self._conns.clear()


def open_account_db(account_dir: str, key: bytes) -> WXWorkDB | None:
    """便捷工厂（兼容旧接口）：单账号目录 + 单密钥。"""
    try:
        user_dir = os.path.basename(os.path.normpath(account_dir))
        return WXWorkDB(
            os.path.dirname(os.path.normpath(account_dir)),
            {user_dir: key.hex()},
        )
    except Exception:
        return None


if __name__ == "__main__":
    # 冒烟测试：真实数据目录 + 密钥，打印可解密账号与最近消息
    keys = load_keys()
    print(f"[*] 密钥文件: {KEYS_FILE}，账号密钥 {len(keys)} 个")
    base = detect_wxwork_dir()
    print(f"[*] 企微数据目录: {base}")
    if not keys or not base:
        print("[-] 缺少密钥或数据目录，无法冒烟")
        raise SystemExit(1)
    db = WXWorkDB(base, keys)
    print(f"[*] 可读账号: {sorted(db._conns.keys())}")
    print(f"[*] health_check: {db.health_check()}")
    msgs = db.get_new_messages(0)
    print(f"[*] 全量消息 {len(msgs)} 条，最近 5 条:")
    for m in msgs[-5:]:
        print(f"    {m['send_time']} [{m['account']}] {m['display_name']}: "
              f"{m['content'][:40]}")
    db.close()