#!/usr/bin/env python3
"""
数据库维护脚本
- 不删除任何历史数据
- 只执行 VACUUM ANALYZE 优化查询性能
- 统计各表数据量和磁盘占用
建议通过 crontab 每周日凌晨执行一次
"""
import os
import sys
import psycopg2
from datetime import datetime

# 数据库配置
DB_CONFIG = {
    'host': os.environ.get('DB_HOST', 'localhost'),
    'port': int(os.environ.get('DB_PORT', 5432)),
    'database': os.environ.get('DB_NAME', 'stock_app'),
    'user': os.environ.get('DB_USER', 'postgres'),
    'password': os.environ.get('DB_PASSWORD', ''),
}


def cleanup():
    """数据库维护：统计数据量 + VACUUM ANALYZE 优化性能"""
    print(f"{'='*60}")
    print(f"[{datetime.now():%Y-%m-%d %H:%M:%S}] 数据库维护开始")
    print(f"  策略：保留全部历史数据，仅执行性能优化")
    print(f"{'='*60}")

    try:
        conn = psycopg2.connect(**DB_CONFIG)
        conn.autocommit = True
        cur = conn.cursor()

        # 1. 统计各表数据量
        print(f"\n  📊 各表数据量统计:")
        cur.execute("""
            SELECT relname as table_name,
                   n_live_tup as row_count,
                   pg_size_pretty(pg_total_relation_size(relid)) as total_size,
                   pg_total_relation_size(relid) as raw_size
            FROM pg_catalog.pg_statio_user_tables
            ORDER BY pg_total_relation_size(relid) DESC
        """)
        total_size = 0
        tables = []
        for name, rows, size, raw in cur.fetchall():
            total_size += raw
            tables.append(name)
            print(f"     {name:>30}: {rows:>12,} 行  {size:>10}")
        print(f"     {'─'*55}")
        print(f"     {'总计':>30}: {'':>12}  {total_size / 1024 / 1024:.0f} MB")

        # 2. VACUUM ANALYZE 优化查询性能
        print(f"\n  🔧 VACUUM ANALYZE (更新统计信息，优化查询)...")
        for table in tables:
            try:
                t0 = datetime.now()
                cur.execute(f"VACUUM ANALYZE {table}")
                elapsed = (datetime.now() - t0).total_seconds()
                if elapsed > 1:
                    print(f"     {table}: {elapsed:.1f}s")
            except Exception as e:
                print(f"     {table}: ⚠ {e}")

        print(f"\n[{datetime.now():%Y-%m-%d %H:%M:%S}] 维护完成 ✅")
        print(f"{'='*60}")

        cur.close()
        conn.close()

    except Exception as e:
        print(f"维护失败: {e}")
        sys.exit(1)


if __name__ == '__main__':
    cleanup()