# -*- coding: utf-8 -*-
"""
test_wal_merge.py - 验证 WAL 帧合并方案（引擎 B 实时性增强）
=============================================================
思路：企微运行时新消息在 message.db-wal（未 checkpoint）。
     WAL 帧数据与主库同算法加密（wxSQLite3 aes128cbc，帧头明文）。
     本脚本：解密主库 → 解密 WAL 帧 → 按页号覆盖合并 → SQLite 打开验证。

用法: python test_wal_merge.py <message.db> <wal_path> <key_hex>
"""
from __future__ import annotations

import os
import sqlite3
import struct
import sys
import tempfile
import time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from wxwork_crypto import (
    decrypt_db_to_file,
    decrypt_page,
    read_page_header,
    SQLITE_FILE_HEADER,
)


def wal_read_frames(wal_path: str, key: bytes):
    """读取 WAL 帧并解密，返回 {pgno: bytes}（仅 salt 匹配的帧，后帧覆盖前帧）。

    返回 (pages, meta)：pages 为页号→解密页映射；meta 含 pgsz/salt/frame_count。
    """
    pages: dict[int, bytes] = {}
    with open(wal_path, "rb") as f:
        hdr = f.read(32)
        if len(hdr) < 32:
            return pages, {}
        magic, ver, pgsz, ckpt, salt1, salt2, _, _ = struct.unpack(">IIIIIIII", hdr)
        meta = {"pgsz": pgsz, "salt1": salt1, "salt2": salt2, "frames": 0, "applied": 0}
        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  # 非当前 WAL 的帧（旧事务/已回滚），跳过
            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
    return pages, meta


def merge_wal_into_db(db_path: str, wal_path: str, key: bytes, out_path: str) -> bool:
    """解密主库 + 合并 WAL 帧 → out_path。返回是否成功。"""
    # 1. 解密主库
    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"]
    # 2. 按页号覆盖
    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)
        # 3. 更新 SQLite 头页数字段（offset 28，4 字节大端）
        if n_pages > 0:
            f.seek(28)
            f.write(struct.pack(">I", n_pages))
    return True


def main():
    if len(sys.argv) < 4:
        print("用法: python test_wal_merge.py <message.db> <wal_path> <key_hex>")
        sys.exit(1)
    db_path, wal_path, key_hex = sys.argv[1], sys.argv[2], sys.argv[3]
    key = bytes.fromhex(key_hex)

    t0 = time.time()
    tmp_fd, tmp_path = tempfile.mkstemp(suffix=".db", prefix="wal_merge_")
    os.close(tmp_fd)
    try:
        ok = merge_wal_into_db(db_path, wal_path, key, tmp_path)
        print(f"[+] 合并完成 ok={ok} 耗时 {time.time()-t0:.2f}s")
        conn = sqlite3.connect(f"file:{tmp_path}?mode=ro", uri=True)
        conn.text_factory = lambda b: b.decode("utf-8", errors="replace")
        try:
            cur = conn.cursor()
            cur.execute("PRAGMA quick_check")
            print("[+] quick_check:", cur.fetchone())
            cur.execute("SELECT COUNT(*) FROM message_table")
            print("[+] message_table 行数:", cur.fetchone()[0])
            cur.execute("SELECT MAX(send_time) FROM message_table")
            mx = cur.fetchone()[0]
            print("[+] MAX(send_time):", mx,
                  time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(mx)) if mx else "-")
            # 与仅主库对比
            cur.execute("SELECT MAX(send_time) FROM message_table WHERE send_time <= (SELECT MAX(send_time) FROM message_table)")
        finally:
            conn.close()
    finally:
        if os.path.exists(tmp_path):
            os.remove(tmp_path)


if __name__ == "__main__":
    main()