"""
第 13 章 · 副本与分布式 · cluster_play.py

演示内容：
  1. 通过 Distributed 表写入 / 查询
  2. 打印 system.clusters（看集群拓扑）
  3. 打印 system.replicas（看副本健康）
  4. 演示 GLOBAL IN / GLOBAL JOIN 的差异
  5. 演示 insert_quorum 强一致写

运行前提：
  - 默认连 127.0.0.1:8123，default 用户，无密码，库 learn_ck
  - 单机情况下 ck_cluster 可能不存在，脚本会先尝试，失败时自动降级到本地表演示

依赖：
  pip install clickhouse-connect
"""

from __future__ import annotations

import os
import sys
import time
from contextlib import contextmanager

try:
    import clickhouse_connect
except ImportError:
    print("请先 pip install clickhouse-connect")
    sys.exit(1)


CK_HOST = os.getenv("CK_HOST", "127.0.0.1")
CK_PORT = int(os.getenv("CK_PORT", "8123"))
CK_USER = os.getenv("CK_USER", "default")
CK_PWD = os.getenv("CK_PWD", "")
CK_DB = os.getenv("CK_DB", "learn_ck")
CK_CLUSTER = os.getenv("CK_CLUSTER", "ck_cluster")


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


def kv(rows, headers=None):
    """打印一个简易表格"""
    if not rows:
        print("  (空结果)")
        return
    if headers:
        widths = [
            max(len(str(h)), max((len(str(r[i])) for r in rows), default=0))
            for i, h in enumerate(headers)
        ]
        print("  " + "  ".join(str(h).ljust(widths[i]) for i, h in enumerate(headers)))
        print("  " + "  ".join("-" * w for w in widths))
        for r in rows:
            print("  " + "  ".join(str(r[i]).ljust(widths[i]) for i in range(len(headers))))
    else:
        for r in rows:
            print("  " + "  ".join(str(c) for c in r))


@contextmanager
def safe_block(title: str):
    """保证示例失败时不中断后续"""
    try:
        yield
    except Exception as exc:  # noqa: BLE001
        print(f"  ⚠ [{title}] 跳过：{exc}")


def main() -> None:
    print(f"连接 {CK_HOST}:{CK_PORT}（user={CK_USER}, db={CK_DB}）...")
    client = clickhouse_connect.get_client(
        host=CK_HOST, port=CK_PORT, username=CK_USER, password=CK_PWD,
        database="default",
    )
    version = client.query("SELECT version()").result_rows[0][0]
    print(f"OK · ClickHouse 版本 {version}")

    # 检测集群是否存在
    cluster_exists = False
    with safe_block("check cluster"):
        rows = client.query(
            f"SELECT count() FROM system.clusters WHERE cluster = '{CK_CLUSTER}'"
        ).result_rows
        cluster_exists = rows and rows[0][0] > 0

    on_cluster = f"ON CLUSTER {CK_CLUSTER}" if cluster_exists else ""
    if not cluster_exists:
        print(f"\n注意：集群 {CK_CLUSTER} 不存在，将以单机/伪分布式模式演示。")
        print("      演示中会自动改用普通 MergeTree 与单分片 Distributed。")

    # ===== 1. 准备 schema =====
    banner("1. 创建演示库与表")
    client.command(f"CREATE DATABASE IF NOT EXISTS {CK_DB} {on_cluster}")
    client.command(f"DROP TABLE IF EXISTS {CK_DB}.events_all   {on_cluster} SYNC")
    client.command(f"DROP TABLE IF EXISTS {CK_DB}.events_local {on_cluster} SYNC")

    if cluster_exists:
        local_engine = (
            "ENGINE = ReplicatedMergeTree("
            "'/clickhouse/tables/{shard}/events_local', '{replica}')"
        )
    else:
        local_engine = "ENGINE = MergeTree()"

    client.command(f"""
        CREATE TABLE {CK_DB}.events_local {on_cluster} (
            event_time DateTime,
            user_id    UInt64,
            event_type LowCardinality(String),
            page       LowCardinality(String),
            properties String,
            country    LowCardinality(String) DEFAULT 'CN'
        ) {local_engine}
        PARTITION BY toYYYYMM(event_time)
        ORDER BY (event_type, user_id, event_time)
    """)

    if cluster_exists:
        client.command(f"""
            CREATE TABLE {CK_DB}.events_all {on_cluster}
            AS {CK_DB}.events_local
            ENGINE = Distributed('{CK_CLUSTER}', '{CK_DB}', 'events_local',
                                 cityHash64(user_id))
        """)
        write_table = f"{CK_DB}.events_all"
    else:
        write_table = f"{CK_DB}.events_local"

    print(f"  写入将走表：{write_table}")

    # ===== 2. 写入演示数据 =====
    banner("2. 写入演示数据")
    rows = [
        ("2026-04-17 09:00:00", 1001, "view",     "home",     '{"src":"app"}',  "CN"),
        ("2026-04-17 09:00:01", 1002, "view",     "home",     '{"src":"web"}',  "US"),
        ("2026-04-17 09:00:02", 1001, "click",    "banner",   '{"id":"b01"}',   "CN"),
        ("2026-04-17 09:00:03", 1003, "view",     "detail",   '{"sku":"A001"}', "JP"),
        ("2026-04-17 09:00:04", 1002, "add_cart", "detail",   '{"sku":"A001"}', "US"),
        ("2026-04-17 09:00:05", 1004, "view",     "home",     '{"src":"app"}',  "CN"),
        ("2026-04-17 09:00:06", 1005, "view",     "home",     '{"src":"app"}',  "CN"),
        ("2026-04-17 09:00:07", 1003, "pay",      "checkout", '{"amount":99}',  "JP"),
        ("2026-04-17 09:00:08", 1006, "view",     "detail",   '{"sku":"B002"}', "DE"),
        ("2026-04-17 09:00:09", 1007, "click",    "banner",   '{"id":"b02"}',   "CN"),
    ]
    cols = ["event_time", "user_id", "event_type", "page", "properties", "country"]
    t0 = time.time()
    client.insert(write_table, rows, column_names=cols)
    print(f"  写入 {len(rows)} 行，耗时 {(time.time()-t0)*1000:.1f} ms")

    # ===== 3. 集群拓扑 =====
    banner("3. system.clusters 集群拓扑")
    res = client.query(f"""
        SELECT cluster, shard_num, replica_num, host_name, port, is_local
        FROM system.clusters
        WHERE cluster = '{CK_CLUSTER}'
        ORDER BY shard_num, replica_num
    """)
    if res.result_rows:
        kv(res.result_rows, headers=["cluster", "shard", "replica", "host", "port", "is_local"])
    else:
        print(f"  集群 {CK_CLUSTER} 不存在（单机模式正常现象）")
        all_clusters = client.query(
            "SELECT DISTINCT cluster FROM system.clusters ORDER BY cluster"
        ).result_rows
        print(f"  本机已知的集群名：{[r[0] for r in all_clusters]}")

    # ===== 4. 副本健康度 =====
    banner("4. system.replicas 副本健康度（仅 Replicated*MergeTree 才有）")
    res = client.query(f"""
        SELECT
            database, table, is_leader, is_readonly,
            queue_size, log_max_index, log_pointer,
            log_max_index - log_pointer AS lag,
            absolute_delay, total_replicas, active_replicas
        FROM system.replicas
        WHERE database = '{CK_DB}'
    """)
    if res.result_rows:
        kv(
            res.result_rows,
            headers=["db", "table", "leader", "readonly", "queue", "log_max",
                     "log_ptr", "lag", "abs_delay", "total", "active"],
        )
    else:
        print("  本表不是 Replicated*MergeTree 引擎（单机模式正常现象）")

    # ===== 5. 查询：fan-out / fan-in =====
    banner("5. 简单聚合 SELECT（在分布式表上跑会自动 fan-out / fan-in）")
    if cluster_exists:
        sql = f"""
            SELECT event_type, uniqExact(user_id) AS uv, count() AS pv
            FROM {CK_DB}.events_all
            GROUP BY event_type
            ORDER BY pv DESC
        """
    else:
        sql = f"""
            SELECT event_type, uniqExact(user_id) AS uv, count() AS pv
            FROM {CK_DB}.events_local
            GROUP BY event_type
            ORDER BY pv DESC
        """
    res = client.query(sql)
    kv(res.result_rows, headers=["event_type", "uv", "pv"])

    # ===== 6. GLOBAL IN 演示（必须有分布式表才有意义） =====
    banner("6. IN vs GLOBAL IN（分布式查询的 N² → N 优化）")
    if cluster_exists:
        # 普通 IN：每个分片会再次发起内层分布式查询
        t0 = time.time()
        r1 = client.query(f"""
            SELECT count() FROM {CK_DB}.events_all
            WHERE user_id IN (
                SELECT user_id FROM {CK_DB}.events_all WHERE event_type = 'pay'
            )
        """).result_rows[0][0]
        t1 = time.time()
        # GLOBAL IN：内层只跑一次，广播给各分片
        r2 = client.query(f"""
            SELECT count() FROM {CK_DB}.events_all
            WHERE user_id GLOBAL IN (
                SELECT user_id FROM {CK_DB}.events_all WHERE event_type = 'pay'
            )
        """).result_rows[0][0]
        t2 = time.time()
        print(f"  普通 IN     →  {r1} 行  耗时 {(t1-t0)*1000:.1f} ms")
        print(f"  GLOBAL IN  →  {r2} 行  耗时 {(t2-t1)*1000:.1f} ms")
        print("  集群越大，差异越明显（GLOBAL IN 复杂度从 N² 降到 N）")
    else:
        print("  单机模式下两者无差异，跳过")

    # ===== 7. insert_quorum 强一致写 =====
    banner("7. insert_quorum 强一致写演示")
    if cluster_exists:
        try:
            client.command(f"""
                INSERT INTO {CK_DB}.events_all
                SETTINGS insert_quorum = 2, insert_quorum_timeout = 10000
                VALUES ('2026-04-17 10:00:00', 9999, 'view', 'home', '{{}}', 'CN')
            """)
            print("  ✓ 强一致写成功（至少 2 副本 ack）")
        except Exception as exc:
            print(f"  ⚠ 强一致写失败：{exc}")
            print("    可能原因：副本数不足、Keeper 不通、超时太短")
    else:
        print("  单机模式下 insert_quorum 无意义，跳过")

    # ===== 8. 看一下数据落到哪些分片 =====
    banner("8. 看每个分片各装了多少行（观察分片键的均衡度）")
    if cluster_exists:
        try:
            res = client.query(f"""
                SELECT hostName() AS host, count() AS rows
                FROM clusterAllReplicas('{CK_CLUSTER}', {CK_DB}.events_local)
                GROUP BY host
                ORDER BY host
            """)
            kv(res.result_rows, headers=["host", "rows"])
        except Exception as exc:
            print(f"  ⚠ 跳过：{exc}")
    else:
        res = client.query(f"SELECT count() FROM {CK_DB}.events_local")
        print(f"  本机总行数：{res.result_rows[0][0]}")

    print("\n演示完成。")


if __name__ == "__main__":
    main()