#!/usr/bin/env python3
"""
物化视图按依赖顺序刷新
用法: python3 refresh_materialized_views.py --db bill_query_test
"""
import argparse
import psycopg2
import time
import sys

from config import DB_CONFIG, MATVIEW_LAYERS


def refresh_all(dbname, verbose=True):
    conn = psycopg2.connect(
        host=DB_CONFIG['host'],
        port=DB_CONFIG['port'],
        dbname=dbname,
        user=DB_CONFIG['user'],
        password=DB_CONFIG['password'] or None,
    )
    conn.autocommit = True
    cur = conn.cursor()

    total = sum(len(layer) for layer in MATVIEW_LAYERS)
    done = 0
    failed = []

    for layer_idx, layer in enumerate(MATVIEW_LAYERS):
        if verbose:
            print(f"\n--- L{layer_idx}: {len(layer)} 个物化视图 ---")
        for mv in layer:
            done += 1
            t0 = time.time()
            try:
                cur.execute(f'REFRESH MATERIALIZED VIEW analytics."{mv}"')
                elapsed = time.time() - t0
                if verbose:
                    print(f"  [{done}/{total}] OK  {mv} ({elapsed:.1f}s)")
            except Exception as e:
                elapsed = time.time() - t0
                err_msg = str(e).strip().split('\n')[0]
                if verbose:
                    print(f"  [{done}/{total}] FAIL {mv} ({elapsed:.1f}s): {err_msg}")
                failed.append((mv, err_msg))
                conn.rollback()

    cur.close()
    conn.close()

    if failed:
        print(f"\n=== 完成: {total - len(failed)}/{total} 成功, {len(failed)} 失败 ===")
        for mv, err in failed:
            print(f"  FAIL {mv}: {err}")
        # 重试失败项（可能因依赖顺序问题，第二轮能成功）
        if len(failed) > 0:
            print(f"\n--- 重试 {len(failed)} 个失败项 ---")
            conn = psycopg2.connect(
                host=DB_CONFIG['host'],
                port=DB_CONFIG['port'],
                dbname=dbname,
                user=DB_CONFIG['user'],
                password=DB_CONFIG['password'] or None,
            )
            conn.autocommit = True
            cur = conn.cursor()
            retry_failed = []
            for mv, _ in failed:
                t0 = time.time()
                try:
                    cur.execute(f'REFRESH MATERIALIZED VIEW analytics."{mv}"')
                    elapsed = time.time() - t0
                    print(f"  OK  {mv} ({elapsed:.1f}s)")
                except Exception as e:
                    err_msg = str(e).strip().split('\n')[0]
                    print(f"  FAIL {mv}: {err_msg}")
                    retry_failed.append(mv)
                conn.rollback()
            cur.close()
            conn.close()
            if retry_failed:
                print(f"\n最终失败: {retry_failed}")
                return 1
            else:
                print(f"\n=== 全部 {total} 个物化视图刷新成功 ===")
                return 0
    else:
        print(f"\n=== 全部 {total} 个物化视图刷新成功 ===")
        return 0


def main():
    parser = argparse.ArgumentParser(description='物化视图按依赖顺序刷新')
    parser.add_argument('--db', default='bill_query_test', help='目标数据库')
    args = parser.parse_args()

    print(f"\n=== 物化视图刷新 ===")
    print(f"数据库: {args.db}")
    rc = refresh_all(args.db)
    sys.exit(rc)


if __name__ == '__main__':
    main()