#!/usr/bin/env python3
"""
第 10 章 · 物化视图 + Projection 查询对比

演示 4 件事:
  1. 直接扫源表 mv_events_raw 算 PV/UV    → 慢 (扫全量明细)
  2. 查 MV 目标表 mv_events_agg + xxxMerge → 快 (扫聚合)
  3. 查 Projection p_hour_agg               → 业务 SQL 不变, 优化器自动选
  4. EXPLAIN PROJECTIONS 看优化器选了谁

前置:
    1. 跑 init.sql 建表
    2. 跑 seed.py --once --total 1000000 (或 --duration 30) 灌点数据

依赖:
    pip install clickhouse-connect
"""
from __future__ import annotations

import argparse
import time
from typing import Tuple

import clickhouse_connect


def time_query(client, sql: str, label: str) -> Tuple[float, list]:
    t0 = time.time()
    rs = client.query(sql)
    cost = (time.time() - t0) * 1000
    print(f"\n=== {label} ===")
    print(f"SQL : {sql.strip()}")
    print(f"耗时: {cost:.1f} ms   返回行数: {len(rs.result_rows)}")
    for row in rs.result_rows[:5]:
        print(f"    {row}")
    if len(rs.result_rows) > 5:
        print(f"    ... ({len(rs.result_rows) - 5} more)")
    return cost, rs.result_rows


def get_table_size(client, table: str) -> str:
    sql = f"""
        SELECT formatReadableSize(sum(bytes_on_disk)) AS sz, sum(rows) AS rows
        FROM system.parts
        WHERE database = 'learn_ck' AND table = '{table}' AND active
    """
    r = client.query(sql).result_rows
    if not r:
        return "(none)"
    sz, rows = r[0]
    return f"{sz} / {rows} rows"


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="")
    args = parser.parse_args()

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

    print("======== 数据规模检查 ========")
    for t in ("mv_events_raw", "mv_events_agg", "mv_events_minute"):
        print(f"  {t:24s} {get_table_size(client, t)}")

    # ---------------- 1. 扫源表 ----------------
    sql_raw = """
        SELECT toStartOfHour(ts) AS hour,
               url,
               count()           AS pv,
               uniq(uid)         AS uv,
               avg(duration)     AS dur_avg,
               quantile(0.95)(duration) AS dur_p95
        FROM learn_ck.mv_events_raw
        WHERE ts >= now() - INTERVAL 1 HOUR
        GROUP BY hour, url
        ORDER BY pv DESC
        LIMIT 10
    """
    cost_raw, _ = time_query(
        client,
        # 强制不命中 projection, 直接扫主数据
        sql_raw + " SETTINGS optimize_use_projections = 0",
        "1. 直接扫明细表 (关闭 Projection)",
    )

    # ---------------- 2. 查 MV 目标表 ----------------
    sql_mv = """
        SELECT hour,
               url,
               sum(pv)                         AS pv,
               uniqMerge(uv_state)             AS uv,
               avgMerge(dur_avg_state)         AS dur_avg,
               quantileMerge(0.95)(dur_p95_state) AS dur_p95
        FROM learn_ck.mv_events_agg
        WHERE hour >= now() - INTERVAL 1 HOUR
        GROUP BY hour, url
        ORDER BY pv DESC
        LIMIT 10
    """
    cost_mv, _ = time_query(client, sql_mv, "2. 查物化视图目标表 (xxxMerge 还原)")

    # ---------------- 3. Projection 自动加速 ----------------
    cost_proj, _ = time_query(
        client,
        sql_raw + " SETTINGS optimize_use_projections = 1",
        "3. 同样的明细 SQL, 让 Projection 自动接管",
    )

    # ---------------- 4. EXPLAIN PROJECTIONS ----------------
    plan = client.query(
        "EXPLAIN PROJECTIONS = 1 " + sql_raw
    ).result_rows
    print("\n=== 4. EXPLAIN PROJECTIONS = 1 看优化器决策 ===")
    for row in plan:
        print(f"    {row[0]}")

    # ---------------- 5. 总结对比 ----------------
    print("\n======== 性能汇总 ========")
    print(f"  扫源表           : {cost_raw:8.1f} ms")
    print(f"  查 MV 聚合表     : {cost_mv:8.1f} ms   ({cost_raw/cost_mv:.1f}x 加速)")
    print(f"  Projection 自动  : {cost_proj:8.1f} ms   ({cost_raw/cost_proj:.1f}x 加速)")
    print("\n结论：")
    print("  - MV: 业务 SQL 要改 (查目标表 + xxxMerge)，但能跨表/链式")
    print("  - Projection: 业务 SQL 零改动, 引擎自动选, 但只能在单表内")


if __name__ == "__main__":
    main()