#!/usr/bin/env python3
"""
第 12 章 · Mutation 实操对比脚本

演示 4 件事:
  1. 灌入 ~100 万行原始数据，观察初始 Part 数
  2. 重型 Mutation:   ALTER TABLE ... UPDATE   →  耗时 + Part 重写
  3. 轻量级 DELETE:   DELETE FROM ...           →  耗时 + 墓碑列
  4. REPLACE PARTITION: staging 表整块替换       →  毫秒级
  + 5. ReplacingMergeTree "软更新" 演示

前置:
    1. clickhouse-client --multiquery < ../init.sql
    2. pip install clickhouse-connect

用法:
    python3 mutation_play.py                  # 全部跑一遍
    python3 mutation_play.py --rows 500000    # 调整数据量
    python3 mutation_play.py --skip-seed      # 不重灌, 复用现有数据
"""
from __future__ import annotations

import argparse
import random
import time
from datetime import date, datetime, timedelta

import clickhouse_connect


def banner(title: str) -> None:
    print("\n" + "=" * 72)
    print(f"  {title}")
    print("=" * 72)


def show_parts(client, table: str) -> dict:
    sql = f"""
        SELECT
            count()                                AS n_parts,
            sum(rows)                              AS rows,
            formatReadableSize(sum(bytes_on_disk)) AS size,
            max(mutations_id)                      AS max_mut
        FROM (
            SELECT name, rows, bytes_on_disk,
                   length(extractAll(name, '_')) - 3 AS mutations_id
            FROM system.parts
            WHERE database='learn_ck' AND table='{table}' AND active
        )
    """
    r = client.query(sql).result_rows
    if not r or r[0][0] == 0:
        return {"n_parts": 0, "rows": 0, "size": "(none)"}
    n, rows, size, mx = r[0]
    return {"n_parts": n, "rows": rows, "size": size, "max_mut": mx}


def wait_mutation(client, table: str, mut_id: str = None, timeout: int = 120) -> dict:
    """阻塞等 Mutation 完成"""
    start = time.time()
    while time.time() - start < timeout:
        cond = f"AND mutation_id='{mut_id}'" if mut_id else ""
        sql = f"""
            SELECT mutation_id, is_done, parts_to_do, latest_fail_reason
            FROM system.mutations
            WHERE database='learn_ck' AND table='{table}' {cond}
            ORDER BY create_time DESC LIMIT 1
        """
        r = client.query(sql).result_rows
        if not r:
            time.sleep(0.5)
            continue
        m_id, done, todo, fail = r[0]
        if done:
            return {"mutation_id": m_id, "done": True, "fail": fail}
        time.sleep(0.5)
    return {"done": False, "timeout": True}


def seed(client, total: int) -> None:
    cols = ["event_date", "event_time", "user_id", "event_type", "payload", "score"]
    today = date.today()
    batch_size = 50_000
    sent = 0
    while sent < total:
        n = min(batch_size, total - sent)
        rows = []
        for _ in range(n):
            d = today - timedelta(days=random.randint(0, 60))
            rows.append((
                d,
                datetime(d.year, d.month, d.day, random.randint(0,23), random.randint(0,59)),
                random.randint(1, 100_000),
                random.choice(["click", "view", "buy", "share"]),
                "P" * random.randint(50, 200),
                random.randint(0, 1000),
            ))
        client.insert("learn_ck.mut_events", rows, column_names=cols)
        sent += n


def demo_mutation(client) -> None:
    banner("Demo 2: 重型 ALTER TABLE UPDATE  → 整 Part 重写")
    before = show_parts(client, "mut_events")
    print(f"  Mutation 前: parts={before['n_parts']}, rows={before['rows']}, size={before['size']}")

    sql = "ALTER TABLE learn_ck.mut_events UPDATE score = score + 1 WHERE event_type = 'click'"
    print(f"\n  执行: {sql}")
    t0 = time.time()
    client.command(sql)
    print(f"  命令立即返回, 耗时 {1000*(time.time()-t0):.1f}ms (异步, 实际还在跑)")

    print("\n  跟踪 system.mutations:")
    res = wait_mutation(client, "mut_events")
    cost = time.time() - t0
    print(f"  Mutation 完成: id={res.get('mutation_id')}  总耗时={cost:.2f}s  fail={res.get('fail')}")

    after = show_parts(client, "mut_events")
    print(f"\n  Mutation 后: parts={after['n_parts']}, rows={after['rows']}, size={after['size']}")
    print(f"  注意 Part 名字带新版本号 (mutation_id), 旧 Part 已 inactive")


def demo_lightweight_delete(client) -> None:
    banner("Demo 3: 轻量级 DELETE FROM (墓碑列) ")
    before = show_parts(client, "mut_events")
    print(f"  DELETE 前: parts={before['n_parts']}, rows={before['rows']}")

    sql = "DELETE FROM learn_ck.mut_events WHERE event_type = 'view' AND score < 100"
    print(f"\n  执行: {sql}")
    t0 = time.time()
    try:
        client.command(sql)
        cost = time.time() - t0
        print(f"  耗时 {1000*cost:.1f}ms (轻量级 DELETE, 写墓碑列, 不重写 Part)")
    except Exception as e:
        print(f"  ⚠ 服务端可能未开启 lightweight delete 或版本 < 22.8: {e}")
        return

    # 看可见行数 (引擎自动加 _row_exists=1 过滤)
    visible = client.query("SELECT count() FROM learn_ck.mut_events").result_rows[0][0]
    after = show_parts(client, "mut_events")
    print(f"\n  DELETE 后:")
    print(f"    SELECT count() (可见行) = {visible}")
    print(f"    system.parts 物理行数 = {after['rows']}  (墓碑还在, 等 Merge 清)")
    print(f"    parts={after['n_parts']}, size={after['size']}")
    print(f"    差值 = {after['rows'] - visible}  (这些是墓碑行)")


def demo_replace_partition(client) -> None:
    banner("Demo 4: REPLACE PARTITION (毫秒级整块替换)")

    # 找一个有数据的分区
    parts = client.query(
        "SELECT DISTINCT partition FROM system.parts "
        "WHERE database='learn_ck' AND table='mut_events' AND active "
        "ORDER BY partition LIMIT 1"
    ).result_rows
    if not parts:
        print("  没数据, 跳过")
        return
    target = parts[0][0]
    print(f"  目标分区: {target}")

    # 1. 清 staging
    client.command("TRUNCATE TABLE learn_ck.mut_events_staging")

    # 2. 在 staging 里准备"修改后"的数据
    sql_stage = f"""
        INSERT INTO learn_ck.mut_events_staging
        SELECT event_date, event_time, user_id,
               'rewritten' AS event_type,
               concat('NEW-', payload) AS payload,
               score * 10 AS score
        FROM learn_ck.mut_events
        WHERE toYYYYMM(event_date) = {target}
    """
    print(f"\n  Step 1: 在 staging 表里准备新版本数据...")
    t0 = time.time()
    client.command(sql_stage)
    print(f"    耗时 {1000*(time.time()-t0):.0f}ms (这部分可以离线慢慢跑)")

    # 3. REPLACE
    print(f"\n  Step 2: ALTER TABLE ... REPLACE PARTITION '{target}' FROM mut_events_staging")
    t0 = time.time()
    client.command(
        f"ALTER TABLE learn_ck.mut_events REPLACE PARTITION '{target}' "
        f"FROM learn_ck.mut_events_staging"
    )
    print(f"    耗时 {1000*(time.time()-t0):.0f}ms  ★ 仅元数据切换!")

    # 验证
    sample = client.query(
        f"SELECT event_type, payload FROM learn_ck.mut_events "
        f"WHERE toYYYYMM(event_date) = {target} LIMIT 1"
    ).result_rows
    print(f"\n  验证: {sample}")


def demo_replacing_mt(client) -> None:
    banner("Demo 5: ReplacingMergeTree 软更新 - INSERT 替代 UPDATE")
    client.command("TRUNCATE TABLE learn_ck.mut_users_dim")

    # 三次"修改"同一个 user_id, 全用 INSERT
    rows = [
        (1, "Alice",   "CN",    datetime.now() - timedelta(days=2)),
        (1, "Alice2",  "CN-SH", datetime.now() - timedelta(days=1)),
        (1, "Alice3",  "CN-BJ", datetime.now()),
    ]
    client.insert("learn_ck.mut_users_dim", rows,
                  column_names=["user_id","nickname","region","updated_at"])

    print("  原始 INSERT 三次, 没 Merge 时:")
    for r in client.query("SELECT * FROM learn_ck.mut_users_dim WHERE user_id=1").result_rows:
        print(f"    {r}")

    print("\n  SELECT FINAL 在线去重:")
    for r in client.query("SELECT * FROM learn_ck.mut_users_dim FINAL WHERE user_id=1").result_rows:
        print(f"    {r}")

    print("\n  推荐: 用 argMax 在线 (不需要 FINAL):")
    sql = """
        SELECT user_id,
               argMax(nickname, updated_at) AS nickname,
               argMax(region,   updated_at) AS region,
               max(updated_at) AS updated_at
        FROM learn_ck.mut_users_dim
        WHERE user_id = 1
        GROUP BY user_id
    """
    for r in client.query(sql).result_rows:
        print(f"    {r}")

    print("\n  强制 OPTIMIZE FINAL 后:")
    client.command("OPTIMIZE TABLE learn_ck.mut_users_dim FINAL")
    for r in client.query("SELECT * FROM learn_ck.mut_users_dim WHERE user_id=1").result_rows:
        print(f"    {r}  ← 物理上只剩最新版本")


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8123)
    parser.add_argument("--user", default="default")
    parser.add_argument("--password", default="")
    parser.add_argument("--rows", type=int, default=300_000)
    parser.add_argument("--skip-seed", action="store_true")
    parser.add_argument("--skip-mutation", action="store_true",
                        help="跳过 demo 2 (重型 Mutation 较慢)")
    args = parser.parse_args()

    client = clickhouse_connect.get_client(
        host=args.host, port=args.port,
        username=args.user, password=args.password,
    )

    if not args.skip_seed:
        banner("Demo 1: 灌入数据")
        client.command("TRUNCATE TABLE learn_ck.mut_events")
        t0 = time.time()
        seed(client, args.rows)
        cost = time.time() - t0
        s = show_parts(client, "mut_events")
        print(f"  灌入 {args.rows} 行  cost={cost:.1f}s  parts={s['n_parts']} size={s['size']}")

    if not args.skip_mutation:
        demo_mutation(client)
    demo_lightweight_delete(client)
    demo_replace_partition(client)
    demo_replacing_mt(client)

    banner("总结")
    print("""
  ┌─────────────────────────────────────────────────────────────┐
  │  操作                          典型耗时    适用场景           │
  ├─────────────────────────────────────────────────────────────┤
  │  ALTER TABLE UPDATE (Mutation) 几分钟~小时  万不得已          │
  │  DELETE FROM (轻量级)           几秒        合规/低频删       │
  │  REPLACE PARTITION              几毫秒      整批修改最佳      │
  │  ReplacingMergeTree + INSERT    毫秒        维度表"修改"      │
  │  CollapsingMergeTree + sign     毫秒        状态翻转          │
  │  TTL DELETE                     0 (自动)    生命周期管理      │
  └─────────────────────────────────────────────────────────────┘
""")


if __name__ == "__main__":
    main()