#!/usr/bin/env python3
"""
测试库与正式库行数对比验证
用法: python3 verify.py --db bill_query_test [--prod bill_query]
"""
import argparse
import psycopg2
import sys

from config import DB_CONFIG


def get_table_counts(dbname):
    conn = psycopg2.connect(
        host=DB_CONFIG['host'],
        port=DB_CONFIG['port'],
        dbname=dbname,
        user=DB_CONFIG['user'],
        password=DB_CONFIG['password'] or None,
    )
    cur = conn.cursor()
    cur.execute("""
        SELECT table_name
        FROM information_schema.tables
        WHERE table_schema = 'public'
          AND table_type = 'BASE TABLE'
        ORDER BY table_name
    """)
    tables = [r[0] for r in cur.fetchall()]

    counts = {}
    for tbl in tables:
        cur.execute(f"SELECT count(*) FROM public.{tbl}")
        counts[tbl] = cur.fetchone()[0]

    # 物化视图行数
    cur.execute("""
        SELECT matviewname
        FROM pg_matviews
        WHERE schemaname = 'analytics'
        ORDER BY matviewname
    """)
    mvs = [r[0] for r in cur.fetchall()]
    mv_counts = {}
    for mv in mvs:
        try:
            cur.execute(f'SELECT count(*) FROM analytics."{mv}"')
            mv_counts[mv] = cur.fetchone()[0]
        except Exception:
            mv_counts[mv] = -1
            conn.rollback()

    cur.close()
    conn.close()
    return counts, mv_counts


def main():
    parser = argparse.ArgumentParser(description='测试库与正式库数据一致性验证')
    parser.add_argument('--db', default='bill_query_test', help='测试库名')
    parser.add_argument('--prod', default='bill_query', help='正式库名')
    args = parser.parse_args()

    print(f"\n=== 数据一致性验证 ===")
    print(f"正式库: {args.prod}")
    print(f"测试库: {args.db}")

    prod_counts, prod_mv = get_table_counts(args.prod)
    test_counts, test_mv = get_table_counts(args.db)

    # 基础表对比
    print(f"\n--- 基础表 ({len(prod_counts)} 张) ---")
    ok = 0
    diff = 0
    for tbl in sorted(prod_counts.keys()):
        p = prod_counts[tbl]
        t = test_counts.get(tbl, 'MISSING')
        flag = 'OK' if t == p else 'DIFF'
        if flag == 'OK':
            ok += 1
        else:
            diff += 1
        if flag != 'OK' or True:  # 显示全部
            print(f"  {flag:4s} {tbl:45s} prod={p:<12s} test={t}")

    print(f"\n  基础表: {ok} OK, {diff} DIFF")

    # 物化视图对比
    print(f"\n--- 物化视图 ({len(prod_mv)} 个) ---")
    mv_ok = 0
    mv_diff = 0
    for mv in sorted(prod_mv.keys()):
        p = prod_mv[mv]
        t = test_mv.get(mv, 'MISSING')
        flag = 'OK' if t == p else 'DIFF'
        if flag == 'OK':
            mv_ok += 1
        else:
            mv_diff += 1
        print(f"  {flag:4s} {mv:45s} prod={p:<12s} test={t}")

    print(f"\n  物化视图: {mv_ok} OK, {mv_diff} DIFF")
    print(f"\n=== 总计: {ok + mv_ok} OK, {diff + mv_diff} DIFF ===")

    if diff + mv_diff > 0:
        sys.exit(1)
    sys.exit(0)


if __name__ == '__main__':
    main()