#!/usr/bin/env python3
"""数据归档脚本

将超过指定天数的旧数据从活跃表移到归档表，减少活跃表体积。
归档表命名规则：原表名 + _archive（如 news_flash → news_flash_archive）

用法：
  python3 archive_data.py --days 90          # 归档90天前的数据
  python3 archive_data.py --days 90 --dry    # 试运行，只统计不执行
  python3 archive_data.py --days 90 --table news_hotlist  # 仅归档指定表
"""
import sys, os, argparse, time
sys.path.insert(0, os.path.dirname(__file__))

from common import MYSQL_CFG
from db_pool import get_conn, put_conn, DBContext

# 需要归档的表列表
ARCHIVE_TABLES = [
    "news_flash",
    "news_hotlist",
    "news_cls", "news_wallstreetcn", "news_sina", "news_36kr",
    "news_eastmoney", "news_jin10", "news_gelonghui", "news_gnews",
    "news_newsapi", "news_fastbull", "news_mktnews",
]


def ensure_archive_table(cursor, table: str):
    """创建归档表（结构与原表相同）"""
    archive = f"{table}_archive"
    cursor.execute(
        "SELECT 1 FROM information_schema.TABLES "
        "WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME = %s",
        (archive,),
    )
    if cursor.fetchone():
        return archive

    cursor.execute(f"CREATE TABLE `{archive}` LIKE `{table}`")
    print(f"  [创建归档表] {archive}")
    return archive


def archive_table(cursor, table: str, cutoff_ms: int, dry: bool = False) -> tuple:
    """归档单张表，返回 (归档条数, 节省空间估算MB)"""
    archive = ensure_archive_table(cursor, table)

    # 统计待归档条数
    cursor.execute(
        f"SELECT COUNT(*) FROM `{table}` WHERE published_at < %s",
        (cutoff_ms,),
    )
    count = cursor.fetchone()[0]
    if count == 0:
        return 0, 0

    if dry:
        print(f"  [试运行] {table}: 待归档 {count} 条")
        return count, 0

    # 迁移数据到归档表
    cursor.execute(
        f"INSERT INTO `{archive}` SELECT * FROM `{table}` WHERE published_at < %s",
        (cutoff_ms,),
    )
    inserted = cursor.rowcount

    # 删除原表旧数据
    cursor.execute(
        f"DELETE FROM `{table}` WHERE published_at < %s",
        (cutoff_ms,),
    )
    deleted = cursor.rowcount

    print(f"  [归档] {table}: 迁移 {inserted} 条, 删除 {deleted} 条")
    return inserted, 0


def main():
    parser = argparse.ArgumentParser(description="FinDataHub 数据归档")
    parser.add_argument("--days", type=int, default=90, help="归档N天前的数据（默认90）")
    parser.add_argument("--dry", action="store_true", help="试运行，只统计不执行")
    parser.add_argument("--table", type=str, default=None, help="仅归档指定表")
    args = parser.parse_args()

    cutoff_ms = int((time.time() - args.days * 86400) * 1000)
    print(f"=== FinDataHub 数据归档 ===")
    print(f"归档 {args.days} 天前的数据 (published_at < {cutoff_ms})")
    if args.dry:
        print("[试运行模式]\n")

    tables = [args.table] if args.table else ARCHIVE_TABLES
    total = 0

    with DBContext() as conn:
        cursor = conn.cursor()
        for table in tables:
            cursor.execute(
                "SELECT 1 FROM information_schema.TABLES "
                "WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME = %s",
                (table,),
            )
            if not cursor.fetchone():
                continue

            print(f">> {table}")
            count, _ = archive_table(cursor, table, cutoff_ms, args.dry)
            total += count

        if not args.dry:
            conn.commit()

        # 优化表空间
        if not args.dry and total > 0:
            print("\n优化表空间...")
            for table in tables:
                cursor.execute(
                    "SELECT 1 FROM information_schema.TABLES "
                    "WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME = %s",
                    (table,),
                )
                if cursor.fetchone():
                    cursor.execute(f"OPTIMIZE TABLE `{table}`")

        cursor.close()

    print(f"\n=== 归档完成: 共 {total} 条 ===")


if __name__ == "__main__":
    main()