#!/usr/bin/env python3
"""
微信聊天 + 朋友圈提取 —— 给 lifelog/photolog 人脉档案页用（里程碑 2/4）。

- 聊天：EnMicroMsg.db（SQLCipher，key e03ce81 + compat=1 + page1024）的 message 表，
  type=1 文本、talker 非 @chatroom → messages.json [{talker,isSend,createTime,content}]。
  注意 sqlcipher .mode json 输出开头有 PRAGMA 的 "ok"，要 raw[raw.index('['):] 跳过，
  否则 app 端 JSONArray 解析失败、消息导不进去（实测踩过）。
- 朋友圈：SnsMicroMsg.db **是明文 SQLite**（不加密！）的 SnsInfo 表，content 是 protobuf，
  正文用「抽含中文最多的可读段」近似提取 → moments.json [{userName,createTime,text}]。

产物 push 到设备 filesDir/wx_enroll/{messages.json,moments.json}，app 内「从微信导入联系人」
的 WeChatEnrollWorker 会 clear+重灌进 wechat_messages / moments 表，按 wxid 关联 Person。

用法：python3 extract-wechat-chat-moments.py /tmp/EnMicroMsg.db /tmp/SnsMicroMsg.db e03ce81 ./out
"""
import sys, os, re, json, subprocess


def sqlcipher_json(db, key, sql):
    script = f'PRAGMA key="{key}";\nPRAGMA cipher_compatibility=1;\nPRAGMA cipher_page_size=1024;\n.mode json\n{sql}\n'
    out = subprocess.run(['sqlcipher', db], input=script, capture_output=True, text=True).stdout
    return json.loads(out[out.index('['):])  # 跳过 PRAGMA 的 "ok"


def extract_messages(en_db, key):
    return sqlcipher_json(en_db, key, """
        SELECT talker, isSend, createTime, content FROM message
        WHERE type=1 AND talker NOT LIKE '%@chatroom' AND content IS NOT NULL AND content != ''
        ORDER BY createTime ASC;""")


def moment_text(blob):
    if not blob:
        return ''
    s = blob.decode('utf-8', 'ignore')
    segs = re.findall(r'[一-鿿A-Za-z0-9，。！？、,.!?#@:：…~（）()【】\-—\s]{2,}', s)

    def score(x):
        x = x.strip()
        if x.startswith(('http', 'wxid', 'mmsns', 'mp.weixin')):
            return -1
        if re.fullmatch(r'[A-Za-z0-9]+', x):
            return -1
        return len(re.findall(r'[一-鿿]', x))
    best = max(segs, key=score, default='')
    return best.strip() if score(best) > 0 else ''


def extract_moments(sns_db):
    import sqlite3  # SnsMicroMsg 是明文
    c = sqlite3.connect(sns_db)
    out = []
    for u, t, content in c.execute("SELECT userName, createTime, content FROM SnsInfo WHERE content IS NOT NULL"):
        txt = moment_text(content)
        if txt:
            ct = int(t)
            ct = ct * 1000 if ct < 1e12 else ct  # 秒→毫秒
            out.append({'userName': u, 'createTime': ct, 'text': txt})
    return out


def main():
    en_db, sns_db, key, out_dir = sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4]
    os.makedirs(out_dir, exist_ok=True)
    msgs = extract_messages(en_db, key)
    json.dump(msgs, open(os.path.join(out_dir, 'messages.json'), 'w'), ensure_ascii=False)
    moments = extract_moments(sns_db)
    json.dump(moments, open(os.path.join(out_dir, 'moments.json'), 'w'), ensure_ascii=False)
    print(f'聊天 {len(msgs)} 条，朋友圈 {len(moments)} 条')


if __name__ == '__main__':
    main()