"""
数据库 K 线读取测试 — 只读，从 TimescaleDB 读取各周期 K 线并打印

用法：
    python db_test.py                    # 使用 env.yaml 中的 host
    python db_test.py --host localhost   # 覆盖 host（如 SSH 隧道后）
"""

import asyncio
import sys

import asyncpg

from common.config import config as app_config

# ── 各周期对应的表/视图 ──
INTERVAL_TABLES: dict[str, str] = {
    "1m": "klines",
    "3m": "klines_3m",
    "5m": "klines_5m",
    "15m": "klines_15m",
    "30m": "klines_30m",
    "1h": "klines_1h",
    "2h": "klines_2h",
    "4h": "klines_4h",
    "6h": "klines_6h",
    "8h": "klines_8h",
    "1d": "klines_1d",
    "1w": "klines_1w",
    "1mon": "klines_1mon",
}

LIMIT = 5


def parse_args() -> str | None:
    """解析命令行参数，返回 host 覆盖值"""
    args = sys.argv[1:]
    host_override = None
    i = 0
    while i < len(args):
        if args[i] == "--host" and i + 1 < len(args):
            host_override = args[i + 1]
            i += 2
        elif args[i].startswith("--host="):
            host_override = args[i].split("=", 1)[1]
            i += 1
        else:
            i += 1
    return host_override


async def main():
    host_override = parse_args()
    db = app_config.db
    if host_override:
        db = db.model_copy(update={"host": host_override})

    dsn = f"postgresql://{db.user}:{db.password}@{db.host}:{db.port}/{db.name}"

    print(f"连接 {db.host}:{db.port}/{db.name} ...")
    conn = await asyncpg.connect(dsn)

    try:
        print()
        print("=" * 85)
        print("  TimescaleDB K 线数据读取测试（只读）")
        print("=" * 85)

        for interval, table in INTERVAL_TABLES.items():
            # 检查表/视图是否存在（含 TimescaleDB 连续聚合视图）
            exists = await conn.fetchval(
                """
                SELECT EXISTS (
                    SELECT 1 FROM pg_matviews WHERE matviewname = $1
                    UNION
                    SELECT 1 FROM pg_tables   WHERE tablename   = $1
                    UNION
                    SELECT 1 FROM pg_views    WHERE viewname    = $1
                )
                """,
                table,
            )

            if not exists:
                print(f"\n  [{interval}] {table} — 不存在，跳过")
                continue

            # 获取该表/视图的实际列名，避免查询不存在的列报错
            columns = await conn.fetch(
                """
                SELECT column_name
                FROM information_schema.columns
                WHERE table_name = $1
                ORDER BY ordinal_position
                """,
                table,
            )
            col_names = {r["column_name"] for r in columns}

            # 选择存在的列进行查询
            select_cols = ["time", "exchange", "symbol", "interval", "open", "high", "low", "close", "volume"]
            if "trade_count" in col_names:
                select_cols.append("trade_count")
            if "is_closed" in col_names:
                select_cols.append("is_closed")

            rows = await conn.fetch(
                f"""
                SELECT {', '.join(select_cols)}
                FROM {table}
                WHERE interval = $1
                ORDER BY time DESC
                LIMIT $2
                """,
                interval,
                LIMIT,
            )

            print(f"\n{'─' * 85}")
            print(f"  [{interval}] {table}  —  {len(rows)} 条")
            print(f"{'─' * 85}")

            if not rows:
                print("  (无数据)")
                continue

            for k in rows:
                t = k["time"].strftime("%Y-%m-%d %H:%M:%S")
                is_closed = k.get("is_closed")
                if is_closed is not None:
                    mark = "✓" if is_closed else "◌"
                else:
                    mark = " "  # 聚合视图无此字段
                trade_count = k.get("trade_count", "-")
                print(
                    f"  {t} [{mark}] {k['symbol']:10s} {k['interval']:4s}"
                    f"  O={k['open']:>12.4f}  H={k['high']:>12.4f}"
                    f"  L={k['low']:>12.4f}  C={k['close']:>12.4f}"
                    f"  V={k['volume']:>10.4f}  trades={trade_count}"
                )

        print(f"\n{'=' * 85}")
        print("  读取完成")
        print("=" * 85)

    finally:
        await conn.close()


if __name__ == "__main__":
    asyncio.run(main())
