#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""merge_constituent.py — 单元3 合并: staging -> constituent_unified (幂等版)。

设计:
- 前置: ``migrate_constituent.py`` 已建好 ``constituent_unified_staging``(全集型, 验证通过)
- 本脚本把 staging 内容复制成正式表 ``constituent_unified``(治幸存者偏差的选股池)

幂等(idempotent, 可重跑, 月度 schtask 安全):
- ``DROP TABLE IF EXISTS constituent_unified``           (旧正式表, 有则删)
- ``CREATE TABLE constituent_unified AS SELECT ... FROM constituent_unified_staging``
- ``CREATE INDEX IF NOT EXISTS idx_constituent_unified`` (index_code, code 加速 get_index_stocks)
- **绝不用 RENAME**(RENAME 只能跑一次, 再跑必崩 —— 早期方案A首次落地用过, 现已废弃)
- **绝不碰 ``bs_index_constituent_old``**(那是 migrate 的事, 这里只管 staging→unified)
- staging 表保留(migrate 下次跑会 DROP+rebuild, 这里不动)

环境变量:
- ``SANGUO_DB``: quant_trading.db 路径, 默认 ``C:\\sanguo_vnpy_v2\\data\\quant_trading.db``(VPS)
"""
import os
import sqlite3


def merge(db_path: str) -> None:
    """从 constituent_unified_staging (幂等) 重建 constituent_unified 正式表。

    Args:
        db_path: quant_trading.db 路径(测试可传 tmp sqlite; 生产读 ``SANGUO_DB``)。

    Raises:
        sqlite3.OperationalError: 若 ``constituent_unified_staging`` 不存在
            (前置 migrate 未跑; 提示先跑 migrate_constituent.py)。
    """
    c = sqlite3.connect(db_path, timeout=60)
    try:
        c.execute("PRAGMA busy_timeout = 60000")

        # 幂等: 先删旧正式表(若存在), 再从 staging CREATE 一张全新的
        # (DROP+CREATE 语义, 不是 RENAME —— 可无限次重跑)
        c.execute("DROP TABLE IF EXISTS constituent_unified")
        c.execute(
            "CREATE TABLE constituent_unified AS "
            "SELECT index_code, code, code_name, source, in_current, was_removed "
            "FROM constituent_unified_staging"
        )
        print("rebuilt constituent_unified from constituent_unified_staging")

        c.execute(
            "CREATE INDEX IF NOT EXISTS idx_constituent_unified "
            "ON constituent_unified(index_code, code)"
        )
        c.commit()

        n_unified = c.execute(
            "SELECT COUNT(*) FROM constituent_unified"
        ).fetchone()[0]
        n_staging = c.execute(
            "SELECT COUNT(*) FROM constituent_unified_staging"
        ).fetchone()[0]
        print(f"constituent_unified rows: {n_unified} (staging: {n_staging})")
    finally:
        c.close()
    print("MERGE DONE")


if __name__ == "__main__":
    _db = os.environ.get(
        "SANGUO_DB", r"C:\sanguo_vnpy_v2\data\quant_trading.db"
    )
    merge(_db)