#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""微信智能客服引擎 v3（wechatauto-replica）

在 v2 固定话术自动回复基础上升级为「大模型上下文对话 + 本机知识库」：

  1. 每个客户（username）维护独立会话历史，持久化到磁盘，跨重启不丢，
     按「最近 N 轮 / 最近 N 秒」窗口取上下文喂给大模型。
  2. 收到消息时检索本机 Obsidian 知识库（字符 n-gram 相似度），把最相关的
     笔记片段注入 system prompt，让回复「结合知识库」。
  3. 通过 OpenAI 兼容接口调用大模型生成自然回复（默认 DeepSeek，可换
     智谱/Kimi/通义等，只要填 base_url + model + api_key）。
  4. 无 api_key 时自动降级：关键词规则 → 知识库片段拼接 → 兜底话术，
     保证引擎不因缺 key 而停摆。

用法（项目目录下，独立 venv）：
    .venv\\Scripts\\python.exe wechat_ai_reply.py

前置：微信 4.x 已登录、桌面未锁屏、管理员权限（提密钥）。
配置：改 ai_config.json（重点是 llm.api_key）。
"""
import json
import os
import re
import sys
import time
import threading
import urllib.request
import urllib.error
import msvcrt

try:
    os.system("chcp 65001 >nul 2>&1")
except Exception:
    pass
try:
    sys.stdout.reconfigure(encoding="utf-8", errors="replace")
except AttributeError:
    pass

from wechatauto.db import WeChatDB, Listener
from wechatauto.guia import WeChatGUI
import voice2text

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
CONFIG_PATH = os.path.join(BASE_DIR, "ai_config.json")
WATERMARK_PATH = os.path.join(BASE_DIR, "data", "watermark.json")


def load_watermark() -> dict:
    """加载监听水位（username -> 已推送的最大 sort_seq），重启不重复推送、不重复回复。"""
    try:
        if os.path.exists(WATERMARK_PATH):
            data = json.loads(open(WATERMARK_PATH, "r", encoding="utf-8").read())
            if isinstance(data, dict):
                return data
    except (OSError, ValueError):
        pass
    return {}


def save_watermark(wm: dict) -> None:
    """把监听水位落盘。"""
    try:
        os.makedirs(os.path.dirname(WATERMARK_PATH), exist_ok=True)
        open(WATERMARK_PATH, "w", encoding="utf-8").write(
            json.dumps(wm, ensure_ascii=False)
        )
    except OSError:
        pass


def load_config() -> dict:
    with open(CONFIG_PATH, "r", encoding="utf-8") as f:
        return json.load(f)


CFG = load_config()

# ==================== 知识库检索（本地 n-gram 相似度） ====================
class KnowledgeBase:
    """读本机 Obsidian vault 的 .md 笔记，按字符 n-gram 重叠度检索。

    与 Unai 本机 vault 的 hashed n-gram 检索思路一致：靠字符重叠而非语义。
    排除 .obsidian 等元数据目录，只读已发布的分类笔记。
    """

    def __init__(self, vault_dir: str):
        self.vault_dir = vault_dir
        self._notes = []  # [(title, path, text)]
        self._load()

    def _load(self):
        if not self.vault_dir or not os.path.isdir(self.vault_dir):
            print(f"[KB] 知识库目录不存在: {self.vault_dir}", flush=True)
            return
        for root, dirs, files in os.walk(self.vault_dir):
            dirs[:] = [d for d in dirs if d not in (".obsidian", ".git")]
            for name in files:
                if not name.endswith(".md"):
                    continue
                path = os.path.join(root, name)
                try:
                    text = open(path, "r", encoding="utf-8", errors="ignore").read()
                except OSError:
                    continue
                if not text.strip():
                    continue
                self._notes.append((name[:-3], path, text))
        print(f"[KB] 已载入 {len(self._notes)} 篇笔记", flush=True)

    @staticmethod
    def _ngrams(s: str, n: int = 3):
        s = re.sub(r"\s+", "", s)
        return {s[i:i + n] for i in range(len(s) - n + 1)} if len(s) >= n else {s}

    def search(self, query: str, top_k: int = 3, max_chars: int = 1500) -> str:
        """返回拼接后的相关笔记片段；无结果返回空串。"""
        if not self._notes or not query:
            return ""
        q_grams = self._ngrams(query)
        scored = []
        for title, path, text in self._notes:
            t_grams = self._ngrams(text, 3)
            inter = len(q_grams & t_grams)
            if inter == 0:
                continue
            # 重叠数 + 少量标题命中加权
            score = inter
            if title and self._ngrams(title) & q_grams:
                score += 10
            scored.append((score, title, text))
        if not scored:
            return ""
        scored.sort(reverse=True, key=lambda x: x[0])
        parts = []
        for score, title, text in scored[:top_k]:
            excerpt = text.strip()
            if len(excerpt) > max_chars:
                excerpt = excerpt[:max_chars] + "…"
            parts.append(f"【{title}】\n{excerpt}")
        return "\n\n".join(parts)


# ==================== 每客户会话记忆 ====================
class ConversationMemory:
    """username -> 独立对话历史，内存缓存 + 磁盘持久化。"""

    def __init__(self, base_dir: str, max_turns: int):
        self.base_dir = base_dir
        self.max_turns = max_turns
        self._cache = {}
        os.makedirs(base_dir, exist_ok=True)

    def _safe_name(self, username: str) -> str:
        return re.sub(r"[^0-9A-Za-z_@-]", "_", username) + ".json"

    def _path(self, username: str) -> str:
        return os.path.join(self.base_dir, self._safe_name(username))

    def load(self, username: str) -> list:
        if username in self._cache:
            return self._cache[username]
        path = self._path(username)
        hist = []
        if os.path.exists(path):
            try:
                hist = json.loads(open(path, "r", encoding="utf-8").read())
            except (OSError, ValueError):
                hist = []
        self._cache[username] = hist
        return hist

    def append(self, username: str, role: str, content: str):
        hist = self.load(username)
        hist.append({"role": role, "content": content, "time": int(time.time())})
        # 只保留最近 max_turns 轮（一轮 = 用户一条 + 助手一条）
        hist = hist[-(self.max_turns * 2):]
        self._cache[username] = hist
        try:
            open(self._path(username), "w", encoding="utf-8").write(
                json.dumps(hist, ensure_ascii=False, indent=1)
            )
        except OSError:
            pass


# ==================== 大模型客户端（OpenAI 兼容，零第三方依赖） ====================
class LLMClient:
    def __init__(self, cfg: dict):
        self.api_key = (cfg.get("api_key") or "").strip()
        self.base_url = (cfg.get("base_url") or "").rstrip("/")
        self.model = cfg.get("model") or "deepseek-chat"
        self.timeout = cfg.get("timeout_sec") or 60
        self.max_retry = cfg.get("max_retry") or 2

    @property
    def available(self) -> bool:
        return bool(self.api_key)

    def chat(self, messages: list) -> str:
        """调用 chat/completions，返回助手文本；失败抛异常。"""
        url = self.base_url + "/chat/completions"
        payload = {
            "model": self.model,
            "messages": messages,
            "temperature": 0.7,
            "stream": False,
        }
        data = json.dumps(payload).encode("utf-8")
        last_err = None
        for attempt in range(self.max_retry + 1):
            req = urllib.request.Request(url, data=data, method="POST")
            req.add_header("Content-Type", "application/json")
            req.add_header("Authorization", "Bearer " + self.api_key)
            try:
                with urllib.request.urlopen(req, timeout=self.timeout) as resp:
                    body = resp.read().decode("utf-8")
                obj = json.loads(body)
                return obj["choices"][0]["message"]["content"].strip()
            except (urllib.error.URLError, urllib.error.HTTPError, KeyError,
                    IndexError, ValueError, OSError) as e:
                last_err = e
                if attempt < self.max_retry:
                    time.sleep(1 + attempt)
        raise RuntimeError(f"LLM 调用失败: {last_err!r}")


# ==================== 引擎主逻辑 ====================
_gui = None
_gui_lock = threading.Lock()
_send_lock = threading.Lock()
_last_reply = {}          # username -> 上次回复时间戳
_last_reply["_db"] = None
_self_wxid = None         # 当前登录账号的 wxid，用于跳过自己发的消息
_handled_seq = {}         # username -> 已处理的最大 sort_seq（消息级去重，双保险）
_sent_recent = {}         # username -> [(content, ts)]：引擎刚发出的回复回声，跳过"自己的消息被读回"
_SELF_SENDER_IDS = {"1"}  # real_sender_id 里"自己"的取值（微信 4.1.13 实测 self=1）
_pending = {}             # username -> {"texts": [..], "timer": Timer|None}：冷却期内累积的消息，到期合并回复
_pending_lock = threading.Lock()


def _recent_sent(username: str, now: float, window: float = 120.0):
    """返回该用户最近 window 秒内发出的回复内容列表，用于回声匹配。"""
    items = [(c, t) for c, t in _sent_recent.get(username, []) if now - t <= window]
    _sent_recent[username] = items
    return [c for c, _ in items]

_kb = KnowledgeBase(CFG["knowledge"].get("vault_dir", ""))
_llm = LLMClient(CFG["llm"])
_mem = ConversationMemory(
    CFG["storage"].get("conversation_dir", "data/conversations"),
    CFG["reply"].get("history_turns", 10),
)


def get_gui():
    global _gui
    with _gui_lock:
        if _gui is None:
            _gui = WeChatGUI()
        return _gui


_BLOCK_PREFIXES = ("gh_",)  # 公众号，不回复
_SYSTEM_USERS = {
    "filehelper", "brandsessionholder", "brandservicesessionholder",
    "notifymessage", "fmessage", "floatbottle", "medianote",
}


def should_reply(username: str) -> bool:
    r = CFG["reply"]
    if not r.get("enabled", True):
        return False
    allowlist = r.get("allowlist") or []
    if allowlist:
        # 白名单模式：只回复名单内的好友（测试期锁定单一联系人用）
        return username in allowlist
    if username in r.get("blocklist", []):
        return False
    if username in _SYSTEM_USERS:
        return False
    if "@" in username:  # 群聊 @chatroom / 服务号 @openim / @weclaw / @placeholder 等
        return False
    if username.startswith(_BLOCK_PREFIXES):
        return False
    return True


def _build_messages(username: str, incoming: str) -> tuple:
    """构造发给大模型的 messages 列表，返回 (messages, kb_hit)。"""
    sys_prompt = CFG["persona"].get("system_prompt", "")
    kb_hit = ""
    kb = _kb.search(incoming, top_k=CFG["knowledge"].get("top_k", 3),
                    max_chars=CFG["knowledge"].get("max_chars_per_note", 1500))
    if kb:
        kb_hit = kb
        sys_prompt += "\n\n【知识库资料】\n" + kb

    messages = [{"role": "system", "content": sys_prompt}]

    # 取该客户历史（按时间窗口过滤）
    hist = _mem.load(username)
    window = CFG["reply"].get("history_window_sec", 86400)
    now = time.time()
    recent = [h for h in hist if now - h.get("time", 0) <= window]
    for h in recent[-(CFG["reply"].get("history_turns", 10) * 2):]:
        if h.get("content"):
            messages.append({"role": h["role"], "content": h["content"]})

    messages.append({"role": "user", "content": incoming})
    return messages, kb_hit


def _fallback_reply(content: str, kb_hit: str) -> str:
    """无 api_key 时的降级：关键词规则 → 知识库片段 → 兜底话术。"""
    rules = CFG["reply"].get("keyword_rules") or {}
    for kw, reply in rules.items():
        if kw and kw in content:
            return reply
    if kb_hit:
        # 知识库命中但无大模型，只回片段开头的摘录 + 引导
        head = kb_hit.strip().split("\n", 1)[0].replace("【", "").replace("】", "")
        return f"关于这个问题，您可以先参考这份资料（{head}）：需要我展开哪一点，我帮您细看。"
    return CFG["reply"].get("fallback_reply",
                            "【自动回复】已收到您的消息，稍后回复您。")


def _do_reply(username: str, content: str, db):
    """生成并发送一次回复（大模型优先，无 key 降级）。"""
    nick = db.get_nickname(username) if db else username
    try:
        if _llm.available:
            messages, kb_hit = _build_messages(username, content)
            reply = _llm.chat(messages)
            mode = "LLM" + ("+KB" if kb_hit else "")
        else:
            kb_hit = _kb.search(content, top_k=CFG["knowledge"].get("top_k", 3),
                                max_chars=CFG["knowledge"].get("max_chars_per_note", 1500))
            reply = _fallback_reply(content, kb_hit)
            mode = "降级(无key)" + ("+KB" if kb_hit else "")
    except Exception as e:  # noqa: BLE001
        reply = CFG["reply"].get("fallback_reply", "【自动回复】已收到您的消息，稍后回复您。")
        mode = f"异常回退:{e!r}"

    # 记录历史（用户 + 助手），再发送（加锁，避免 timer 线程与 worker 线程并发写）
    with _send_lock:
        _mem.append(username, "user", content)
        _mem.append(username, "assistant", reply)
        _sent_recent.setdefault(username, []).append((reply, time.time()))
        try:
            r = get_gui().send_msg(reply, who=username, verify=True)
            print(f"[智能回复|{mode}] {username} ({nick}) <- {reply} | ok={r.is_success}",
                  flush=True)
        except Exception as e:  # noqa: BLE001
            print(f"[发送失败] {username}: {e!r}", flush=True)


def _flush_pending(username: str):
    """冷却期结束后，把累积的多条消息合并成一次回复发出。"""
    with _pending_lock:
        p = _pending.get(username)
        if not p or not p["texts"]:
            return
        texts = p["texts"]
        p["texts"] = []
        p["timer"] = None
    content = " ".join(texts).strip()
    if not content:
        return
    db = _last_reply["_db"]
    _last_reply[username] = time.time()
    _do_reply(username, content, db)


def on_message(msg: dict, lst: Listener):
    username = msg.get("username", "")
    sender_id = str(msg.get("sender_id", ""))
    content = str(msg.get("content", "")).strip()
    sort_seq = msg.get("sort_seq", 0)
    mtype = msg.get("type", "")
    local_id = msg.get("local_id")

    # 跳过自己发的（避免回复自己的回复 → 死循环）。
    now = time.time()

    # ---- 跳过"自己"的消息（三层防护，杜绝自回复死循环）----
    # 1) sender_id 是真实 wxid 字符串时
    if _self_wxid and sender_id == str(_self_wxid):
        return
    # 2) real_sender_id 数值 == 自己（微信 4.1.13 实测 self=1）
    if sender_id in _SELF_SENDER_IDS:
        return
    # 3) 内容回声：刚发出去的回复又被监听读回，直接丢弃（不依赖 sender_id 语义，最稳一层）
    if content and content in _recent_sent(username, now):
        return

    if not should_reply(username) or not content:
        return

    # 消息级去重：sort_seq 不前进则说明是同一条消息被重复回调，直接丢弃
    last_seq = _handled_seq.get(username, -1)
    if sort_seq and sort_seq <= last_seq:
        return
    _handled_seq[username] = sort_seq

    db = _last_reply["_db"]

    # 语音消息：先转文字再回复（语音 content 是 "[语音]" 占位，真实内容在音频里）
    if mtype == "语音" and local_id:
        text = voice2text.transcribe_voice(username, local_id, db)
        if text:
            content = text
            print(f"[语音转文字] {username} -> {text}", flush=True)
        else:
            print(f"[语音转文字] {username} 转写失败，跳过", flush=True)
            return

    # 冷却合并：冷却期内的新消息不丢弃，累积后统一回复（修复"只回第一句、后面不回"）。
    cooldown = CFG["reply"].get("cooldown_sec", 5)
    if now - _last_reply.get(username, 0) < cooldown:
        with _pending_lock:
            p = _pending.setdefault(username, {"texts": [], "timer": None})
            p["texts"].append(content)
            if p["timer"] is not None:
                p["timer"].cancel()
            p["timer"] = threading.Timer(cooldown, _flush_pending, args=(username,))
            p["timer"].daemon = True
            p["timer"].start()
        return

    _last_reply[username] = now
    _do_reply(username, content, db)


def main():
    global _self_wxid
    # 单实例锁：防止多个引擎进程并发读写同一个解密缓存导致 database malformed
    _lock_path = os.path.join(BASE_DIR, "engine.lock")
    _lock_fh = open(_lock_path, "wb")
    try:
        msvcrt.locking(_lock_fh.fileno(), msvcrt.LK_NBLCK, 1)
    except OSError:
        print("已有引擎实例在运行，本进程退出（避免并发写坏数据库缓存）。", flush=True)
        _lock_fh.close()
        sys.exit(0)

    db = WeChatDB()
    info = db.get_self_info()
    nick = info.get("nick_name") or info.get("username")
    _last_reply["_db"] = db
    _self_wxid = info.get("username") or getattr(db, "wxid", None)
    print(f"已接管微信: {nick} ({info.get('username')})", flush=True)
    print(f"自身 wxid: {_self_wxid}（跳过自己发的消息，防自回复循环）", flush=True)
    print(f"大模型: {'已配置 ' + _llm.model if _llm.available else '未配置(降级模式)'} "
          f"| base_url={_llm.base_url}", flush=True)
    print(f"知识库: {CFG['knowledge'].get('vault_dir','')} (已载入 {len(_kb._notes)} 篇)",
          flush=True)
    _allow = CFG['reply'].get('allowlist') or []
    _scope = f"白名单 {len(_allow)} 人" if _allow else "所有好友(排除公众号/群聊/服务号)"
    print(f"自动回复: {'开' if CFG['reply'].get('enabled') else '关'} | "
          f"回复范围: {_scope} | "
          f"防抖 {CFG['reply'].get('cooldown_sec')}s", flush=True)

    lst = Listener(db, interval=3.0, watermark=load_watermark())
    # 监听所有会话（add_all + discover 自动发现新会话）。should_reply 里排除
    # 公众号/群聊/服务号/系统账号；若配置 allowlist 则只回名单内好友。
    lst.add_all(on_message, discover=True)
    print(f"智能客服监听已启动（{_scope}），Ctrl+C 退出...",
          flush=True)

    lst.start()
    _save_every = 0
    try:
        while True:
            time.sleep(5)
            _save_every += 5
            if _save_every >= 30:
                _save_every = 0
                save_watermark(lst.watermark)
    except KeyboardInterrupt:
        pass
    finally:
        save_watermark(lst.watermark)
        lst.stop()
        print("已停止监听。", flush=True)


if __name__ == "__main__":
    main()