#!/usr/bin/env python3
"""
统计每个联系人的聊天记录数量和最近/最早消息时间
输出: wechat_plain/chat_stats.json  { username: {count, last_time, first_time} }
"""
import sqlite3, hashlib, json, os, time, sys

DB_DIR     = '/Users/yibo/Documents/wechat_plain'
CONTACT_DB = f'{DB_DIR}/contact_contact.db'
MSG_DBS    = [f'{DB_DIR}/message_message_{i}.db' for i in range(11)
              if os.path.exists(f'{DB_DIR}/message_message_{i}.db')]
OUT_FILE   = f'{DB_DIR}/chat_stats.json'

def md5(s):
    return hashlib.md5(s.encode()).hexdigest()

# ── 加载联系人 username，构建 md5 → username 反查表 ─────────────────────────
print('[1/3] 加载联系人...')
md5_to_user = {}
try:
    con = sqlite3.connect(CONTACT_DB)
    for (u,) in con.execute("SELECT username FROM contact WHERE delete_flag=0"):
        md5_to_user[md5(u)] = u
    con.close()
    print(f'      {len(md5_to_user)} 个联系人')
except Exception as e:
    print(f'[!] 联系人读取失败: {e}')
    sys.exit(1)

# ── 遍历所有消息表，统计 count/last_time/first_time ─────────────────────────
print('[2/3] 统计消息表...')
stats = {}   # username (or md5 hash) → {count, last_time, first_time}
total_tables = 0
t0 = time.time()

for db_path in MSG_DBS:
    db_name = os.path.basename(db_path)
    try:
        con = sqlite3.connect(db_path)
        tables = [r[0] for r in con.execute(
            "SELECT name FROM sqlite_master WHERE type='table' AND name LIKE 'Msg_%'"
        )]
        print(f'      {db_name}: {len(tables)} 张表', flush=True)
        for tbl in tables:
            h = tbl[4:]
            key = md5_to_user.get(h, h)   # prefer real username, fallback to md5
            try:
                row = con.execute(
                    f'SELECT count(*), max(create_time), min(create_time) FROM "{tbl}"'
                ).fetchone()
                count, last_t, first_t = row
                if count and count > 0:
                    prev = stats.get(key, {'count': 0, 'last_time': 0, 'first_time': 0})
                    stats[key] = {
                        'count':      prev['count'] + count,
                        'last_time':  max(prev['last_time'],  last_t  or 0),
                        'first_time': min(prev['first_time'] or (first_t or 0), first_t or 0),
                    }
            except Exception:
                pass
            total_tables += 1
        con.close()
    except Exception as e:
        print(f'[!] {db_name}: {e}')

elapsed = time.time() - t0
print(f'      {total_tables} 张表, {len(stats)} 个聊天对象  ({elapsed:.1f}s)')

# ── 写入 JSON ─────────────────────────────────────────────────────────────────
print(f'[3/3] 写入 {OUT_FILE}...')
with open(OUT_FILE, 'w') as f:
    json.dump(stats, f, ensure_ascii=False)

size_kb = os.path.getsize(OUT_FILE) // 1024
print(f'      完成！{len(stats)} 条记录  ({size_kb} KB)')
print(f'\n前 10 名（消息最多）:')
top = sorted(stats.items(), key=lambda x: x[1]['count'], reverse=True)[:10]
for key, v in top:
    print(f'  {key[:40]:<40}  {v["count"]:>6} 条')