# -*- coding: utf-8 -*-
"""用「备份 ZIP 里的真实旧库」验证：迁移新列后旧数据必须 100% 原样保留（只读对比）。

用法: python tests/check_db_migration.py <backup.zip>
"""
import os
import sys
import json
import shutil
import sqlite3
import tempfile
import zipfile

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

ZIP = sys.argv[1] if len(sys.argv) > 1 else r"D:/haocode_backup_20260916_1102.zip"
OUT = os.path.join(tempfile.gettempdir(), "haocode_from_backup.db")
if os.path.exists(OUT):
    os.remove(OUT)

with zipfile.ZipFile(ZIP) as z:
    name = [n for n in z.namelist()
            if os.path.basename(n) == "chat_history.db"][0]
    with z.open(name) as src, open(OUT, "wb") as dst:
        shutil.copyfileobj(src, dst)
print("从备份取出:", name, os.path.getsize(OUT), "bytes")


def fp(path):
    c = sqlite3.connect(path)
    c.row_factory = sqlite3.Row
    cols = [r[1] for r in c.execute("PRAGMA table_info(messages)")]
    d = {
        "messages_rows": c.execute("SELECT COUNT(*) FROM messages").fetchone()[0],
        "sessions_rows": c.execute("SELECT COUNT(*) FROM sessions").fetchone()[0],
        "msg_sum": list(c.execute(
            "SELECT COUNT(*), SUM(LENGTH(COALESCE(content,''))), "
            "SUM(LENGTH(COALESCE(timeline,''))) FROM messages").fetchone()[:]),
        "parents": [tuple(r) for r in c.execute(
            "SELECT id, parent_id FROM messages ORDER BY id")],
        "leaves": [tuple(r) for r in c.execute(
            "SELECT id, current_leaf_msg_id FROM sessions ORDER BY id")],
        "branches": sorted((r["parent_id"], r["n"]) for r in c.execute(
            "SELECT parent_id, COUNT(*) n FROM messages "
            "WHERE parent_id IS NOT NULL GROUP BY parent_id HAVING n>1")),
        "roles": sorted((r["role"], r["n"]) for r in c.execute(
            "SELECT role, COUNT(*) n FROM messages GROUP BY role")),
        "has_new_cols": all(x in cols for x in ("stop_reason", "error_message")),
    }
    c.close()
    return d


before = fp(OUT)
print("备份库(迁移前): 消息=%d 会话=%d 分支点=%d 含新列=%s" % (
    before["messages_rows"], before["sessions_rows"],
    len(before["branches"]), before["has_new_cols"]))

import core.db_manager as _dbm  # noqa: E402
_dbm._DEFAULT_DB = OUT
d = _dbm.DBManager(db_path=OUT)          # ← 触发迁移
after = fp(OUT)
print("迁移后        : 消息=%d 会话=%d 分支点=%d 含新列=%s" % (
    after["messages_rows"], after["sessions_rows"],
    len(after["branches"]), after["has_new_cols"]))
print()
keys = ("messages_rows", "sessions_rows", "msg_sum", "parents",
        "leaves", "branches", "roles")
for k in keys:
    print(("  OK   " if before[k] == after[k] else "  DIFF ") + "%-14s" % k)
allok = all(before[k] == after[k] for k in keys)
print()
print("每个会话叶子:", [(r[0][:16], (r[1] or "")[:16]) for r in after["leaves"][:4]], "...")
print(">>> 结论:", "✅ 旧库数据 100% 原样保留（只多两个空列）"
      if allok else "❌ 有改动！")
os.remove(OUT)
sys.exit(0 if allok else 1)