#!/usr/bin/env python3
"""
微信联系人头像提取 —— 给 lifelog/photolog 的人脸锚点 enrollment 备数据。

流程（在 Mac 上跑，需要 `brew install sqlcipher`）：
  1. adb 把设备上的 EnMicroMsg.db 拉到本地（需 root）。
  2. 用微信的 SQLCipher 配方解密：key + cipher_compatibility=1 + page 1024。
  3. join rcontact ⨝ img_flag，筛真人（verifyFlag=0，排除群/公众号/企业微信/系统号），
     img_flag.reserved1 = 高清头像 URL（wx.qlogo.cn，公开可拉）。
  4. 并行下载头像 → avatars/{username}.jpg。
  5. （可选）insightface 估"有脸"覆盖率，导出 oneface 清单。

产物（敏感 PII，勿入库）放 lifelog-recorder/_wx_enroll/：contacts.json / avatars/ / oneface_usernames.json
然后 push 到设备 filesDir/wx_enroll/，app「从微信导入联系人」用端上 FaceEngine 建锚点。

  DB key / 路径因设备而异，见 android-wechat-tracker-sqlcipher.js（key 'e03ce81'，
  路径 /data/data/com.tencent.mm/MicroMsg/<hash>/EnMicroMsg.db）。

用法：
  python3 extract-wechat-avatars.py /tmp/EnMicroMsg.db e03ce81 ./_wx_enroll
"""
import sys, os, json, subprocess, urllib.request, concurrent.futures

REAL_PERSON_SQL = """
SELECT r.username AS username,
       COALESCE(NULLIF(r.conRemark,''), NULLIF(r.nickname,''), r.username) AS name,
       r.nickname AS nickname, r.conRemark AS remark, r.alias AS alias,
       r.pyInitial AS pyInitial, r.quanPin AS quanPin,
       f.reserved1 AS avatar
FROM rcontact r JOIN img_flag f ON f.username=r.username
WHERE r.verifyFlag=0
  AND r.username NOT LIKE '%@chatroom'
  AND r.username NOT LIKE 'gh_%'
  AND r.username NOT LIKE '%@openim'
  AND r.username NOT IN ('weixin','filehelper','fmessage','medianote','floatbottle',
        'qmessage','qqmail','newsapp','officialaccounts','helper_entry','exmail_tool',
        'mphelper','brandsessionholder')
  AND f.reserved1 LIKE 'http%';
"""


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


def download(items, out_dir):
    os.makedirs(out_dir, exist_ok=True)

    def dl(it):
        dest = os.path.join(out_dir, it['username'] + '.jpg')
        if os.path.exists(dest) and os.path.getsize(dest) > 0:
            return True
        try:
            req = urllib.request.Request(it['avatar'], headers={'User-Agent': 'Mozilla/5.0'})
            b = urllib.request.urlopen(req, timeout=12).read()
            if len(b) < 200:
                return False
            open(dest, 'wb').write(b)
            return True
        except Exception:
            return False

    with concurrent.futures.ThreadPoolExecutor(max_workers=24) as ex:
        return sum(1 for r in ex.map(dl, items) if r)


def main():
    db_path, key, out_root = sys.argv[1], sys.argv[2], sys.argv[3]
    os.makedirs(out_root, exist_ok=True)
    items = decrypt_and_query(db_path, key)
    json.dump(items, open(os.path.join(out_root, 'contacts.json'), 'w'), ensure_ascii=False)
    print(f'真人联系人(有头像): {len(items)}')
    n = download(items, os.path.join(out_root, 'avatars'))
    print(f'头像下载: {n}')
    print('下一步：（可选）insightface 估覆盖率导出 oneface_usernames.json，'
          '再 push 到设备 filesDir/wx_enroll/，app 内「从微信导入联系人」。')


if __name__ == '__main__':
    main()