#!/usr/bin/env python3
"""增量补全 hk_songs_test（或任意目标表）中缺失的 bpm_class 字段。

从源库 hk_music_record_state.bpm 取主版本录音的 BPM 值，按以下规则映射：
    bpm = 0 或 NULL  →  NULL（跳过）
    bpm < 80         →  1（慢歌）
    80 ≤ bpm ≤ 120   →  2（中速）
    bpm > 120        →  3（快歌）

用法:
    python backfill_bpm_class.py                     # 默认更新 hk_songs_test
    python backfill_bpm_class.py --table hk_songs    # 更新其他表
    python backfill_bpm_class.py --dry-run           # 只打印，不写库
    python backfill_bpm_class.py --batch-size 500    # 每批处理条数（默认200）
"""

from __future__ import annotations

import argparse
import logging
import os

import pymysql
from dotenv import load_dotenv
from tqdm import tqdm

load_dotenv(os.path.join(os.path.dirname(__file__), ".env"))

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(message)s",
    datefmt="%H:%M:%S",
)
logger = logging.getLogger(__name__)

TARGET_DB_CONFIG = {
    "host": os.getenv("TARGET_DB_HOST"),
    "port": int(os.getenv("TARGET_DB_PORT", 3306)),
    "user": os.getenv("TARGET_DB_USER"),
    "password": os.getenv("TARGET_DB_PASSWORD"),
    "database": os.getenv("TARGET_DB_NAME"),
    "charset": "utf8mb4",
    "cursorclass": pymysql.cursors.DictCursor,
}

SOURCE_DB_CONFIG = {
    "host": os.getenv("SOURCE_DB_HOST"),
    "port": int(os.getenv("SOURCE_DB_PORT", 3306)),
    "user": os.getenv("SOURCE_DB_USER"),
    "password": os.getenv("SOURCE_DB_PASSWORD"),
    "database": os.getenv("SOURCE_DB_NAME"),
    "charset": "utf8mb4",
    "cursorclass": pymysql.cursors.DictCursor,
}

# 主版本优先：同一首歌取 is_main_version=1 的录音，没有则取 min(record_id)
BPM_QUERY = """
SELECT sp.id AS song_id, rs.bpm
FROM hk_song_platform sp
INNER JOIN (
    SELECT song_id,
           COALESCE(
               MIN(CASE WHEN is_main_version = 1 THEN record_id END),
               MIN(record_id)
           ) AS record_id
    FROM hk_song_and_record
    GROUP BY song_id
) sr ON sr.song_id = sp.id
LEFT JOIN hk_music_record r
    ON r.id = sr.record_id AND r.deleted = 0x00
LEFT JOIN hk_music_record_state rs
    ON rs.record_id = r.id AND rs.deleted = 0x00
WHERE sp.id IN ({placeholders})
"""


def bpm_to_class(bpm_str: str | None) -> int | None:
    if not bpm_str:
        return None
    try:
        bpm = int(float(bpm_str))
    except (ValueError, TypeError):
        return None
    if bpm <= 0:
        return None
    if bpm < 80:
        return 1
    if bpm <= 120:
        return 2
    return 3


def fetch_bpm_map(src_conn, song_ids: list[str]) -> dict[str, int | None]:
    """从源库批量查 bpm，返回 {source_song_id: bpm_class}。"""
    if not song_ids:
        return {}
    placeholders = ",".join(["%s"] * len(song_ids))
    sql = BPM_QUERY.format(placeholders=placeholders)
    with src_conn.cursor() as cur:
        cur.execute(sql, song_ids)
        rows = cur.fetchall()
    return {str(r["song_id"]): bpm_to_class(r["bpm"]) for r in rows}


def backfill(table: str, batch_size: int, dry_run: bool) -> None:
    tgt = pymysql.connect(**TARGET_DB_CONFIG)
    src = pymysql.connect(**SOURCE_DB_CONFIG)

    try:
        # 查需要补全的记录总数
        with tgt.cursor() as cur:
            cur.execute(
                f"SELECT COUNT(*) AS cnt FROM `{table}` "
                f"WHERE bpm_class IS NULL AND source_song_id IS NOT NULL AND deleted = '0'"
            )
            total = cur.fetchone()["cnt"]

        logger.info(f"待补全 bpm_class 的记录数: {total}，表: {table}")
        if total == 0:
            logger.info("无需更新，退出。")
            return

        updated = no_bpm = 0
        last_id = 0  # 用 id > last_id 游标代替 OFFSET，避免跳过无bpm记录

        with tqdm(total=total, unit="条") as bar:
            while True:
                with tgt.cursor() as cur:
                    cur.execute(
                        f"SELECT id, source_song_id FROM `{table}` "
                        f"WHERE bpm_class IS NULL AND source_song_id IS NOT NULL AND deleted = '0' "
                        f"AND id > %s ORDER BY id LIMIT %s",
                        (last_id, batch_size),
                    )
                    rows = cur.fetchall()

                if not rows:
                    break

                last_id = rows[-1]["id"]
                song_ids = [r["source_song_id"] for r in rows]
                bpm_map = fetch_bpm_map(src, song_ids)

                for row in rows:
                    sid = row["source_song_id"]
                    bpm_class = bpm_map.get(str(sid))

                    if bpm_class is None:
                        no_bpm += 1
                        bar.update(1)
                        continue

                    if dry_run:
                        logger.debug(f"[DRY-RUN] id={row['id']} source_song_id={sid} -> bpm_class={bpm_class}")
                        updated += 1
                        bar.update(1)
                        continue

                    with tgt.cursor() as cur:
                        cur.execute(
                            f"UPDATE `{table}` SET bpm_class = %s WHERE id = %s",
                            (bpm_class, row["id"]),
                        )
                    tgt.commit()
                    updated += 1
                    bar.update(1)

        prefix = "[DRY-RUN] " if dry_run else ""
        print(f"\n{'=' * 50}")
        print(f"{prefix}bpm_class 补全完成")
        print(f"  成功更新: {updated} 条")
        print(f"  源库无bpm: {no_bpm} 条（保持 NULL）")
        print(f"  合计处理: {updated + no_bpm} 条（共 {total} 条待补全）")
        print(f"{'=' * 50}")

    finally:
        tgt.close()
        src.close()


def main() -> None:
    parser = argparse.ArgumentParser(description="增量补全 bpm_class 字段")
    parser.add_argument("--table", default="hk_songs_test", help="目标表名（默认 hk_songs_test）")
    parser.add_argument("--batch-size", type=int, default=200, help="每批处理条数（默认 200）")
    parser.add_argument("--dry-run", action="store_true", help="只打印计划，不写库")
    args = parser.parse_args()
    backfill(args.table, args.batch_size, args.dry_run)


if __name__ == "__main__":
    main()