#!/usr/bin/env python3
"""数据库索引优化脚本

优化内容：
1. news_flash 主表添加复合索引（importance+is_pushed+is_dup, published_at+importance）
2. 各分表添加复合索引（importance+is_dup, importance+is_pushed）
3. 输出当前索引状态供验证
"""
import sys, os
sys.path.insert(0, os.path.dirname(__file__))

from common import MYSQL_CFG, KX_HOME, KXS_HOME, BUNDLED_KXS_HOME
from db_pool import get_conn, DBContext


# 分表列表（与 dal.py VALID_TABLES 对齐）
SOURCE_TABLES = [
    "news_cls", "news_wallstreetcn", "news_sina", "news_36kr",
    "news_eastmoney", "news_jin10", "news_gnews", "news_gelonghui",
    "news_newsapi", "news_hotlist", "news_fastbull", "news_mktnews",
]

# news_flash 主表新增索引
MAIN_TABLE_INDEXES = [
    # 复合索引：待推送查询 (importance >= 3 AND is_pushed = 0 AND is_duplicate = 0)
    ("idx_importance_pushed", "news_flash", "(importance, is_pushed, is_duplicate)"),
    # 复合索引：按时间+热度查询
    ("idx_published_importance", "news_flash", "(published_at, importance)"),
]

# 分表新增索引
SOURCE_TABLE_INDEXES = [
    # 复合索引：待评分查询 (importance = 0 AND is_dup = 0)
    ("idx_imp_dup", "(importance, is_dup)"),
    # 复合索引：待推送查询 (importance >= 3 AND is_pushed = 0 AND is_dup = 0)
    ("idx_imp_pushed_dup", "(importance, is_pushed, is_dup)"),
]


def check_index_exists(cursor, table: str, index_name: str) -> bool:
    cursor.execute(
        "SELECT 1 FROM information_schema.STATISTICS "
        "WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME = %s AND INDEX_NAME = %s LIMIT 1",
        (table, index_name),
    )
    return cursor.fetchone() is not None


def add_index_safe(cursor, index_name: str, table: str, columns: str):
    if check_index_exists(cursor, table, index_name):
        print(f"  [skip] {table}.{index_name} 已存在")
        return
    sql = f"ALTER TABLE `{table}` ADD INDEX `{index_name}` {columns}"
    print(f"  [add]  {table}.{index_name} {columns}")
    cursor.execute(sql)


def optimize_main_table(cursor):
    print("[1/2] 优化 news_flash 主表索引...")
    for index_name, table, columns in MAIN_TABLE_INDEXES:
        add_index_safe(cursor, index_name, table, columns)


def optimize_source_tables(cursor):
    print("[2/2] 优化分表索引...")
    for table in SOURCE_TABLES:
        cursor.execute(
            "SELECT 1 FROM information_schema.TABLES "
            "WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME = %s LIMIT 1",
            (table,),
        )
        if not cursor.fetchone():
            print(f"  [skip] {table} 不存在")
            continue

        print(f"  >> {table}")
        for index_name, columns in SOURCE_TABLE_INDEXES:
            add_index_safe(cursor, index_name, table, columns)


def show_index_status(cursor):
    print("\n=== 索引状态 ===")
    # 主表
    cursor.execute("SHOW INDEX FROM news_flash")
    rows = cursor.fetchall()
    print(f"\nnews_flash ({len(rows)} 个索引):")
    seen = set()
    for r in rows:
        name = r[2]
        if name not in seen:
            seen.add(name)
            print(f"  - {name}")

    # 分表
    for table in SOURCE_TABLES:
        cursor.execute(
            "SELECT 1 FROM information_schema.TABLES "
            "WHERE TABLE_SCHEMA = DATABASE() AND TABLE_NAME = %s LIMIT 1",
            (table,),
        )
        if not cursor.fetchone():
            continue
        cursor.execute(f"SHOW INDEX FROM `{table}`")
        rows = cursor.fetchall()
        seen = set()
        for r in rows:
            name = r[2]
            if name not in seen:
                seen.add(name)
        print(f"  {table}: {', '.join(sorted(seen))}")


def main():
    print("=== FinDataHub 数据库索引优化 ===\n")
    with DBContext() as conn:
        cursor = conn.cursor()
        optimize_main_table(cursor)
        optimize_source_tables(cursor)
        conn.commit()
        show_index_status(cursor)
        cursor.close()
    print("\n=== 索引优化完成 ===")


if __name__ == "__main__":
    main()