"""
第 14 章 · 集成生态 · integration_play.py

演示内容：
  1. mysql() 表函数：跨库查询
  2. postgresql() 表函数：跨库查询
  3. s3() 表函数：读公开 Parquet（无需 AK/SK 的公开桶）
  4. url() 表函数：读公开 CSV
  5. remote() 表函数：跨集群 ad-hoc

凭证 / 地址全部从【环境变量】读取，避免硬编码：
  CK_HOST              = 127.0.0.1   (CK HTTP 端口 8123)
  MYSQL_HOST / DB / USER / PWD / TABLE
  PG_HOST   / DB / USER / PWD / TABLE / SCHEMA
  S3_URL               = https://<bucket>/path/*.parquet
  S3_AK   / S3_SK      = (可选)
  REMOTE_CK            = other-ck:9000
  REMOTE_USER / REMOTE_PWD

不设置环境变量也能跑：会自动跳过对应章节，仅演示 ClickHouse 端的 SQL 模板。

依赖：
  pip install clickhouse-connect
"""

from __future__ import annotations

import os
import sys
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")


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


@contextmanager
def section(title: str):
    banner(title)
    try:
        yield
    except Exception as exc:  # noqa: BLE001
        print(f"  ⚠ 跳过：{type(exc).__name__}: {exc}")


def main() -> None:
    print(f"连接 {CK_HOST}:{CK_PORT}（user={CK_USER}）...")
    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}")
    client.command(f"CREATE DATABASE IF NOT EXISTS {CK_DB}")

    # =====================================================================
    # 1) mysql() 表函数
    # =====================================================================
    with section("1) mysql() 表函数（跨库查询 MySQL）"):
        host = os.getenv("MYSQL_HOST")
        if not host:
            print("  未设置 MYSQL_HOST 环境变量，仅打印 SQL 模板：")
            print("""
  SELECT user_id, count() AS cnt
  FROM mysql('mysql-host:3306','db','table','user','pwd')
  WHERE created_at >= today() - 7
  GROUP BY user_id
  ORDER BY cnt DESC
  LIMIT 5;
            """.strip())
        else:
            db = os.getenv("MYSQL_DB", "shop")
            tbl = os.getenv("MYSQL_TABLE", "users")
            user = os.getenv("MYSQL_USER", "root")
            pwd = os.getenv("MYSQL_PWD", "")
            sql = f"""
                SELECT count() FROM mysql('{host}', '{db}', '{tbl}', '{user}', '{pwd}')
            """
            res = client.query(sql).result_rows
            print(f"  MySQL {host}/{db}.{tbl} 行数 = {res[0][0]}")

    # =====================================================================
    # 2) postgresql() 表函数
    # =====================================================================
    with section("2) postgresql() 表函数（跨库查询 PostgreSQL）"):
        host = os.getenv("PG_HOST")
        if not host:
            print("  未设置 PG_HOST 环境变量，仅打印 SQL 模板：")
            print("""
  SELECT * FROM postgresql(
      'pg-host:5432','db','table','user','pwd','public'
  ) WHERE status = 'paid' LIMIT 10;
            """.strip())
        else:
            db = os.getenv("PG_DB", "shop")
            tbl = os.getenv("PG_TABLE", "orders")
            user = os.getenv("PG_USER", "postgres")
            pwd = os.getenv("PG_PWD", "")
            schema = os.getenv("PG_SCHEMA", "public")
            sql = f"""
                SELECT count() FROM postgresql(
                    '{host}','{db}','{tbl}','{user}','{pwd}','{schema}'
                )
            """
            res = client.query(sql).result_rows
            print(f"  PG {host}/{db}.{schema}.{tbl} 行数 = {res[0][0]}")

    # =====================================================================
    # 3) url() 表函数：读公开 CSV
    # =====================================================================
    with section("3) url() 表函数：把网络上的 CSV 当成 CK 表查"):
        # 用 CK 官方公开数据集做演示（如果环境能联网）
        url = os.getenv(
            "DEMO_URL",
            "https://datasets.clickhouse.com/hits/tsv/hits_v1.tsv.gz"
        )
        # 仅打印 SQL，不真去 download（百 MB 级别）
        print(f"  示例 SQL（不实际执行重型下载）：")
        print(f"""
  SELECT count() FROM url(
    '{url}',
    'TSV',
    'WatchID UInt64, ...'
  );
        """.strip())

        # 真正跑一个轻量级的 demo：读一个非常小的 CSV
        small_url = os.getenv(
            "SMALL_CSV",
            "https://gist.githubusercontent.com/curran/a08a1080b88344b0c8a7/raw/iris.csv"
        )
        try:
            sql = f"""
                SELECT count() FROM url(
                    '{small_url}', 'CSVWithNames'
                )
            """
            res = client.query(sql).result_rows
            print(f"  小 CSV 行数 = {res[0][0]} (URL: {small_url})")
        except Exception as exc:  # noqa: BLE001
            print(f"  ⚠ 网络受限，跳过：{exc}")

    # =====================================================================
    # 4) s3() 表函数
    # =====================================================================
    with section("4) s3() 表函数：直接读 S3 上的 Parquet/CSV"):
        s3_url = os.getenv("S3_URL")
        if not s3_url:
            print("  未设置 S3_URL 环境变量，仅打印 SQL 模板：")
            print("""
  -- 读
  SELECT count() FROM s3(
      'https://my-bucket.s3.amazonaws.com/events/*.parquet',
      'Parquet',
      'AKIA...', 'SECRET...'
  );

  -- 写（归档）
  INSERT INTO FUNCTION s3(
      'https://my-bucket.s3.amazonaws.com/archive/2026.parquet',
      'Parquet'
  )
  SELECT * FROM learn_ck.events_local
  WHERE event_time >= '2026-01-01';
            """.strip())
        else:
            ak = os.getenv("S3_AK", "")
            sk = os.getenv("S3_SK", "")
            fmt = os.getenv("S3_FORMAT", "Parquet")
            if ak and sk:
                sql = f"SELECT count() FROM s3('{s3_url}', '{fmt}', '{ak}', '{sk}')"
            else:
                sql = f"SELECT count() FROM s3('{s3_url}', '{fmt}')"
            res = client.query(sql).result_rows
            print(f"  S3 {s3_url} 行数 = {res[0][0]}")

    # =====================================================================
    # 5) remote() 表函数
    # =====================================================================
    with section("5) remote() 表函数：跨集群 / 跨节点 ad-hoc"):
        rh = os.getenv("REMOTE_CK")
        if not rh:
            print("  未设置 REMOTE_CK 环境变量，仅打印 SQL 模板：")
            print("""
  SELECT count() FROM remote(
      'other-ck:9000', learn_ck.events_local, 'user', 'pwd'
  );

  -- 跨集群迁移：
  INSERT INTO new_db.events
  SELECT * FROM remote('old_cluster_name', old_db.events);
            """.strip())
        else:
            ru = os.getenv("REMOTE_USER", "default")
            rp = os.getenv("REMOTE_PWD", "")
            res = client.query(
                f"SELECT count() FROM remote('{rh}', system.tables, '{ru}', '{rp}')"
            ).result_rows
            print(f"  远端 system.tables 行数 = {res[0][0]}")

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


if __name__ == "__main__":
    main()