#!/usr/bin/env python3
"""
第 18 章 · chdb 嵌入式 ClickHouse 演示
=========================================

干什么:
    1) 用 chdb 直接在 Python 进程里跑 ClickHouse SQL, 无需 server.
    2) 演示三种典型形态:
         (a) 纯 SQL 计算 (numbers/range)
         (b) 直接读取本地 CSV / Parquet 文件
         (c) Pandas DataFrame ↔ chdb 互通 (DataFrame as table)
    3) 演示 Session 模式: 在内存里建表, 多次查询复用 + 持久化到本地目录.

依赖:
    pip install chdb pandas pyarrow

运行:
    python chdb_demo.py
"""
from __future__ import annotations

import os
import sys
import tempfile
import textwrap
from pathlib import Path

try:
    import chdb
    from chdb import session as chs
except ImportError:
    print("[ERROR] 请先 pip install chdb", file=sys.stderr)
    sys.exit(1)

try:
    import pandas as pd
except ImportError:
    print("[ERROR] 请先 pip install pandas pyarrow", file=sys.stderr)
    sys.exit(1)


def banner(title: str) -> None:
    print()
    print("=" * 60)
    print(f"  {title}")
    print("=" * 60)


# ------------------------------------------------------------
# Demo 1: 一行 SQL 算 1 ~ 1e6 的和 (纯 stateless)
# ------------------------------------------------------------
def demo_stateless() -> None:
    banner("Demo 1 · stateless: chdb.query 直接出结果")
    res = chdb.query(
        "SELECT count() AS cnt, sum(number) AS total FROM numbers(1000000)",
        "PrettyCompact",
    )
    print(res)


# ------------------------------------------------------------
# Demo 2: 用 chdb 读 CSV / 写 Parquet
# ------------------------------------------------------------
def demo_files(workdir: Path) -> Path:
    banner("Demo 2 · file(): 把本地 CSV 当成表")

    csv_path = workdir / "sales.csv"
    csv_path.write_text(textwrap.dedent("""\
        order_id,country,amount,ts
        1,CN,120.50,2026-04-01 10:00:00
        2,US, 88.00,2026-04-01 10:05:00
        3,CN,330.10,2026-04-01 10:30:00
        4,JP, 50.00,2026-04-01 11:00:00
        5,CN,210.00,2026-04-01 11:30:00
        6,US,400.00,2026-04-01 12:00:00
        7,JP,160.00,2026-04-01 12:10:00
        8,DE, 75.00,2026-04-01 12:30:00
    """), encoding="utf-8")

    sql = f"""
        SELECT country,
               count()           AS orders,
               round(sum(amount), 2) AS gmv
        FROM file('{csv_path}', CSVWithNames)
        GROUP BY country
        ORDER BY gmv DESC
        FORMAT PrettyCompact
    """
    print(chdb.query(sql))

    parquet_path = workdir / "sales.parquet"
    chdb.query(f"""
        SELECT * FROM file('{csv_path}', CSVWithNames)
        INTO OUTFILE '{parquet_path}'
        FORMAT Parquet
    """)
    print(f"[OK] 已写出 Parquet: {parquet_path} "
          f"({parquet_path.stat().st_size} bytes)")
    return parquet_path


# ------------------------------------------------------------
# Demo 3: Pandas DataFrame 互通
#   chdb 0.x 起支持: 把 DataFrame 当成虚拟表, 用 SQL 直接查
# ------------------------------------------------------------
def demo_dataframe() -> None:
    banner("Demo 3 · DataFrame 互通: SQL on Pandas")

    df = pd.DataFrame({
        "user_id": [1, 1, 2, 2, 3, 3, 3],
        "event":   ["view", "buy", "view", "view", "view", "buy", "buy"],
        "amount":  [0, 99.0, 0, 0, 0, 12.5, 8.0],
    })
    print("[原始 DataFrame]")
    print(df)
    print()

    sql = """
        SELECT user_id,
               countIf(event = 'view') AS views,
               countIf(event = 'buy')  AS buys,
               round(sumIf(amount, event = 'buy'), 2) AS spend
        FROM Python(df)
        GROUP BY user_id
        ORDER BY user_id
        FORMAT PrettyCompact
    """
    try:
        print(chdb.query(sql))
    except Exception as e:  # 老版本 chdb 可能不支持 Python() 表函数
        print(f"[WARN] 当前 chdb 版本不支持 Python() 表函数: {e}")
        print("[FALLBACK] 用 csv 中转")
        with tempfile.NamedTemporaryFile(
            mode="w", suffix=".csv", delete=False
        ) as f:
            df.to_csv(f.name, index=False)
            tmp = f.name
        sql2 = f"""
            SELECT user_id,
                   countIf(event = 'view') AS views,
                   countIf(event = 'buy')  AS buys,
                   round(sumIf(amount, event = 'buy'), 2) AS spend
            FROM file('{tmp}', CSVWithNames)
            GROUP BY user_id
            ORDER BY user_id
            FORMAT PrettyCompact
        """
        print(chdb.query(sql2))
        os.unlink(tmp)


# ------------------------------------------------------------
# Demo 4: Session 模式 + 持久化 (本地目录就是数据库)
# ------------------------------------------------------------
def demo_session(workdir: Path, parquet: Path) -> None:
    banner("Demo 4 · Session 模式: 本地目录 = 数据库, 多次查询复用")

    db_dir = workdir / "chdb_data"
    db_dir.mkdir(exist_ok=True)

    sess = chs.Session(str(db_dir))
    sess.query("CREATE DATABASE IF NOT EXISTS demo")
    sess.query("""
        CREATE TABLE IF NOT EXISTS demo.sales
        ENGINE = MergeTree
        ORDER BY ts
        AS SELECT * FROM file('{path}', Parquet)
    """.format(path=parquet))

    print("[demo.sales 行数]")
    print(sess.query("SELECT count() FROM demo.sales", "PrettyCompact"))

    print("[按国家汇总]")
    print(sess.query("""
        SELECT country, count() AS orders, round(sum(amount), 2) AS gmv
        FROM demo.sales
        GROUP BY country
        ORDER BY gmv DESC
        FORMAT PrettyCompact
    """))

    sess.close()
    print(f"[OK] 关闭 session. 数据落地于 {db_dir} (下次直接 reopen 即可)")


def main() -> None:
    print("chdb version:", getattr(chdb, "__version__", "unknown"))

    with tempfile.TemporaryDirectory(prefix="chdb_demo_") as tmp:
        workdir = Path(tmp)
        demo_stateless()
        parquet = demo_files(workdir)
        demo_dataframe()
        demo_session(workdir, parquet)

    banner("DONE")
    print("以上 4 个 demo 均不需要 ClickHouse server, 全部跑在 Python 进程内.")


if __name__ == "__main__":
    main()