"""
basic_query.py —— 第 2 章配套代码：5 种 ClickHouse 常用查询模式

涵盖：
  1) 简单 count / SELECT
  2) 分组聚合 + LIMIT BY（每组取前 N）
  3) WITH FILL 自动补齐时间序列
  4) 用 SETTINGS 子句临时覆盖参数（max_threads / max_memory_usage）
  5) 多种 FORMAT 输出（PrettyCompact / JSONEachRow / CSVWithNames / Vertical）

运行前提：
  - ClickHouse 已在 127.0.0.1:8123 启动
  - 已经执行过 02_client_basic/init.sql（建好 learn_ck.events 表并灌好 1 万行）
  - pip install -r requirements.txt

运行：
  python 02_client_basic/code/basic_query.py
"""

from __future__ import annotations

import sys

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


CK = dict(host="127.0.0.1", port=8123, username="default", password="", database="learn_ck")


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


def demo_1_count(client) -> None:
    hr("DEMO 1: 最简 count + SELECT，验证表是否就绪")
    rows = client.query("SELECT count() FROM events").result_rows
    print(f"events 表共 {rows[0][0]:,} 行")
    rows = client.query("SELECT * FROM events ORDER BY ts LIMIT 3").result_rows
    print("\n前 3 行示例：")
    for r in rows:
        print(" ", r)


def demo_2_groupby_limit_by(client) -> None:
    hr("DEMO 2: GROUP BY + LIMIT BY —— 每个 country 取 PV 前 3 的 URL")
    sql = """
        SELECT country, url, count() AS pv
        FROM events
        GROUP BY country, url
        ORDER BY country, pv DESC
        LIMIT 3 BY country
    """
    rows = client.query(sql).result_rows
    cur_country = None
    for country, url, pv in rows:
        if country != cur_country:
            print(f"\n[{country}]")
            cur_country = country
        print(f"  {url:<20}  pv={pv}")


def demo_3_with_fill(client) -> None:
    hr("DEMO 3: WITH FILL —— 自动补齐缺失的小时（即使该小时没有事件）")
    sql = """
        SELECT toStartOfHour(ts) AS h, count() AS pv
        FROM events
        WHERE ts BETWEEN toDateTime('2026-04-01 00:00:00')
                     AND toDateTime('2026-04-01 12:00:00')
        GROUP BY h
        ORDER BY h
        WITH FILL
            FROM toDateTime('2026-04-01 00:00:00')
            TO   toDateTime('2026-04-01 12:00:00')
            STEP toIntervalHour(1)
    """
    rows = client.query(sql).result_rows
    print("小时              | pv")
    print("-" * 32)
    for h, pv in rows:
        bar = "█" * min(int(pv / 5), 30)
        print(f"  {h}  {pv:>4}  {bar}")


def demo_4_settings(client) -> None:
    hr("DEMO 4: SETTINGS 子句 —— 临时覆盖 max_threads")
    sql = """
        SELECT browser, count() AS pv, avg(response_ms) AS avg_ms
        FROM events
        GROUP BY browser
        ORDER BY pv DESC
        SETTINGS max_threads = 2, max_memory_usage = 1000000000
    """
    res = client.query(sql)
    print(f"{'browser':<12} {'pv':>8} {'avg_ms':>10}")
    print("-" * 32)
    for browser, pv, avg_ms in res.result_rows:
        print(f"{browser:<12} {pv:>8} {avg_ms:>10.1f}")
    print(f"\n服务端汇报：扫描 {res.summary.get('read_rows','?')} 行；"
          f"耗时 {int(res.summary.get('elapsed_ns', 0)) / 1e6:.2f} ms")


def demo_5_formats(client) -> None:
    hr("DEMO 5: FORMAT 全家桶 —— 同一份数据 4 种格式输出")
    base_sql = "SELECT country, count() AS pv FROM events GROUP BY country ORDER BY pv DESC LIMIT 3"

    print("\n--- FORMAT PrettyCompact（CLI 默认，人眼好看）---")
    print(client.raw_query(f"{base_sql} FORMAT PrettyCompact").decode())

    print("--- FORMAT JSONEachRow（流式 JSON Lines）---")
    print(client.raw_query(f"{base_sql} FORMAT JSONEachRow").decode())

    print("--- FORMAT CSVWithNames（导出 / Excel 友好）---")
    print(client.raw_query(f"{base_sql} FORMAT CSVWithNames").decode())

    print("--- FORMAT Vertical（字段太多时调试用）---")
    print(client.raw_query(f"SELECT * FROM events ORDER BY ts LIMIT 1 FORMAT Vertical").decode())


def main() -> None:
    try:
        client = clickhouse_connect.get_client(**CK)
    except Exception as e:
        sys.stderr.write(f"❌ 连接失败：{e}\n")
        sys.exit(1)

    if not client.query("EXISTS TABLE learn_ck.events").result_rows[0][0]:
        sys.stderr.write(
            "❌ learn_ck.events 表不存在，请先执行：\n"
            "   clickhouse-client --multiquery < 02_client_basic/init.sql\n"
        )
        sys.exit(1)

    demo_1_count(client)
    demo_2_groupby_limit_by(client)
    demo_3_with_fill(client)
    demo_4_settings(client)
    demo_5_formats(client)

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


if __name__ == "__main__":
    main()