"""
hello_ck.py —— ClickHouse 教程第 1 章配套代码：第一次握手 ClickHouse

本脚本会做 4 件事：
  1) 连接本机 8123 HTTP 端口；
  2) 打印服务器版本（SELECT version()）；
  3) 创建演示库 learn_ck（IF NOT EXISTS）；
  4) 在 learn_ck 库下建一张 1000 行的演示表 events，灌入随机数据，
     然后跑一条聚合查询，体验「列存 + 向量化」一秒返回的快感。

运行前提：
  - ClickHouse Server 已在 127.0.0.1:8123 启动（默认账户 default，无密码）；
    如果还没起来，进入仓库根目录执行：docker compose up -d
  - Python 3.8+；pip install -r requirements.txt（核心是 clickhouse-connect）

运行：
  python 01_intro/code/hello_ck.py

预期输出（节选）：
  ✅ 连接成功，ClickHouse 版本：24.8.x
  ✅ learn_ck 库已就绪
  ✅ events 表已重建并灌入 1000 行
  📊 按 country 分组的 PV 排行：
       BJ  → 332
       SH  → 245
       ...
  ⏱ 聚合查询耗时：3.21 ms
"""

from __future__ import annotations

import random
import sys
import time
from datetime import datetime, timedelta

try:
    import clickhouse_connect
except ImportError:
    sys.stderr.write(
        "❌ 缺少依赖 clickhouse-connect，请先执行：\n"
        "   pip install -r requirements.txt\n"
    )
    sys.exit(1)


CK_HOST = "127.0.0.1"
CK_PORT = 8123
CK_USER = "default"
CK_PASSWORD = ""
CK_DB = "learn_ck"


def connect():
    """建立到 ClickHouse 的 HTTP 连接。"""
    return clickhouse_connect.get_client(
        host=CK_HOST,
        port=CK_PORT,
        username=CK_USER,
        password=CK_PASSWORD,
        # 默认 connect_timeout=10s, send_receive_timeout=300s，足够本教程使用
    )


def step_version(client) -> str:
    """步骤 1：握手 + 打印版本。"""
    rows = client.query("SELECT version()").result_rows
    version = rows[0][0]
    print(f"✅ 连接成功，ClickHouse 版本：{version}")
    return version


def step_create_db(client) -> None:
    """步骤 2：创建演示库 learn_ck（容器化部署时已自动建好，这里幂等执行一次）。"""
    client.command(f"CREATE DATABASE IF NOT EXISTS {CK_DB}")
    print(f"✅ {CK_DB} 库已就绪")


def step_create_table(client) -> None:
    """步骤 3：建一张演示表 events。

    设计要点：
      - country 用 LowCardinality(String) —— 第 3 章重点，少量重复值的字典编码；
      - ts     用 DateTime CODEC(DoubleDelta, ZSTD)) —— 时序列经典编码；
      - 表引擎选用 MergeTree —— 第 4、5 章会详细讲；
      - ORDER BY (country, ts) 把数据按国家 + 时间排好，未来按 country 过滤超快。
    """
    client.command(f"DROP TABLE IF EXISTS {CK_DB}.events")
    client.command(
        f"""
        CREATE TABLE {CK_DB}.events
        (
            event_id    UInt64,
            user_id     UInt32,
            country     LowCardinality(String),
            url         String                CODEC(ZSTD(3)),
            bytes       UInt32                CODEC(T64, LZ4),
            ts          DateTime              CODEC(DoubleDelta, ZSTD(3))
        )
        ENGINE = MergeTree
        ORDER BY (country, ts)
        """
    )


def step_seed(client, n_rows: int = 1000) -> None:
    """步骤 4：灌入 n_rows 行随机数据。

    采用 client.insert() 批量写入 —— 这才是 ClickHouse 的正确写入姿势：
    单条 INSERT 是反模式，千万行也建议一批写完（详见第 7 章）。
    """
    countries = ["BJ", "SH", "GZ", "SZ", "HZ", "CD", "WH", "XA"]
    base = datetime(2026, 4, 17, 0, 0, 0)

    rows = []
    for i in range(n_rows):
        rows.append(
            (
                i + 1,                                        # event_id
                random.randint(1, 1000),                      # user_id
                random.choice(countries),                     # country
                f"/page/{random.randint(1, 50)}",            # url
                random.randint(200, 200000),                  # bytes
                base + timedelta(seconds=i * 7),              # ts
            )
        )

    client.insert(
        f"{CK_DB}.events",
        rows,
        column_names=["event_id", "user_id", "country", "url", "bytes", "ts"],
    )
    print(f"✅ events 表已重建并灌入 {n_rows} 行")


def step_aggregate(client) -> None:
    """步骤 5：跑一条经典 OLAP 聚合查询，体验「列存 + 向量化」的速度。

    SQL 语义：按国家分组，统计访问次数 / 平均字节数，按 PV 倒序。
    生产里这种查询一般在亿级表上跑，本例只灌了 1000 行，主要演示流程。
    """
    sql = f"""
        SELECT
            country,
            count()                AS pv,
            avg(bytes)             AS avg_bytes,
            quantile(0.9)(bytes)   AS p90_bytes
        FROM {CK_DB}.events
        GROUP BY country
        ORDER BY pv DESC
    """

    t0 = time.perf_counter()
    result = client.query(sql)
    elapsed_ms = (time.perf_counter() - t0) * 1000.0

    print("\n📊 按 country 分组的 PV 排行：")
    print(f"   {'country':<10} {'pv':>8} {'avg_bytes':>12} {'p90_bytes':>12}")
    print(f"   {'-' * 44}")
    for country, pv, avg_bytes, p90 in result.result_rows:
        print(f"   {country:<10} {pv:>8} {avg_bytes:>12.1f} {p90:>12.1f}")

    print(f"\n⏱ 聚合查询耗时：{elapsed_ms:.2f} ms")


def main() -> None:
    print("=" * 60)
    print(f"Hello ClickHouse  →  http://{CK_HOST}:{CK_PORT}")
    print("=" * 60)

    try:
        client = connect()
    except Exception as e:
        sys.stderr.write(
            f"❌ 连不上 {CK_HOST}:{CK_PORT}：{e}\n"
            "   请确认 ClickHouse 已在该地址启动；\n"
            "   或在仓库根目录执行：docker compose up -d\n"
        )
        sys.exit(1)

    step_version(client)
    step_create_db(client)
    step_create_table(client)
    step_seed(client, n_rows=1000)
    step_aggregate(client)

    print("\n🎉 第 1 章实操完成。打开 01_intro/demo.html 在浏览器里看可视化演示。")


if __name__ == "__main__":
    main()