"""验证 Hummingbot 那 1030ms 的来源：Bitget candle1m 的 snapshot / update 序列。

Hummingbot 的 bitget_perpetual candles feed 里有这么一句：

    if data and data.get("data") and data["action"] == "update":

action == "snapshot" 的消息被整条丢弃。若新 K 线的首条推送恰是 snapshot，
Hummingbot 就必须等到下一条 update 才知道换根了，代价约等于一个推送间隔。

本脚本直接连原始 WS，对每条消息记录 action、K 线时间戳、到达时刻，
然后针对每次换根回答两件事：
  1. 首条带新时间戳的消息，action 是什么
  2. 若是 snapshot，到首条 update 之间隔了多久（= Hummingbot 白等的时间）

    .venv/bin/python research/live/probe_ws_action.py --minutes 6
"""
from __future__ import annotations

import argparse
import asyncio
import json
import time
from collections import defaultdict
from pathlib import Path

WSS = "wss://ws.bitget.com/v2/ws/public"
SYMS = ("BTCUSDT", "ETHUSDT", "SOLUSDT")
OUT = Path(__file__).resolve().parents[1] / "out" / "probe_ws_action.csv"


async def run(minutes: int) -> None:
    import csv

    import aiohttp

    payload = {"op": "subscribe",
               "args": [{"instType": "USDT-FUTURES", "channel": "candle1m",
                         "instId": s} for s in SYMS]}

    n_rows = 0
    # 每个币记录：当前时间戳、该时间戳下已见过的 action 序列
    cur: dict[str, int] = {}
    seen: dict[str, list] = defaultdict(list)
    rollovers: list[dict] = []
    cnt: dict = defaultdict(lambda: defaultdict(int))

    # 边收边写：进程若被杀，已采到的样本仍在盘上
    OUT.parent.mkdir(parents=True, exist_ok=True)
    fh = OUT.open("w", newline="")
    w = csv.DictWriter(fh, fieldnames=["t_ms", "sym", "action", "kline_ts"])
    w.writeheader()

    print(f"[原始 WS 探针] {SYMS} · candle1m · 跑 {minutes} 分钟", flush=True)
    deadline = time.time() + minutes * 60
    try:
        while time.time() < deadline:
            try:
                async with aiohttp.ClientSession() as sess, \
                        sess.ws_connect(WSS, heartbeat=20) as ws:
                    await ws.send_str(json.dumps(payload))
                    while time.time() < deadline:
                        msg = await ws.receive(timeout=30)
                        # 断开后 receive() 会立刻返回 CLOSED；这里若只 continue
                        # 就成了无等待空转，必须跳出去重连
                        if msg.type is not aiohttp.WSMsgType.TEXT:
                            print(f"  非文本消息 {msg.type}，重连", flush=True)
                            break
                        raw = msg.data
                        if raw == "pong":
                            continue
                        try:
                            d = json.loads(raw)
                        except Exception:
                            continue
                        if "data" not in d or "arg" not in d:
                            if d.get("event"):
                                print(f"  事件: {d}", flush=True)
                            continue

                        now_ms = int(time.time() * 1000)
                        sym = d["arg"]["instId"]
                        action = d.get("action")
                        kts = int(d["data"][0][0])
                        w.writerow({"t_ms": now_ms, "sym": sym,
                                    "action": action, "kline_ts": kts})
                        n_rows += 1
                        cnt[sym][action] += 1
                        if n_rows % 50 == 0:
                            fh.flush()

                        prev = cur.get(sym)
                        if prev is None:
                            cur[sym] = kts
                            seen[sym] = [(action, now_ms)]
                            continue
                        if kts > prev:
                            cur[sym] = kts
                            seen[sym] = [(action, now_ms)]
                            rollovers.append({"sym": sym, "kline_ts": kts,
                                              "first_action": action,
                                              "first_ms": now_ms})
                            print(f"  {sym} 换根 -> 首条 action={action}",
                                  flush=True)
                        else:
                            seen[sym].append((action, now_ms))
                            # 若首条是 snapshot，找该时间戳下首条 update 的延后量
                            for r in rollovers:
                                if (r["sym"] == sym and r["kline_ts"] == kts
                                        and r["first_action"] == "snapshot"
                                        and "gap_to_update_ms" not in r
                                        and action == "update"):
                                    r["gap_to_update_ms"] = now_ms - r["first_ms"]
                                    print(f"    {sym} snapshot->update 间隔 "
                                          f"{r['gap_to_update_ms']}ms", flush=True)
            except asyncio.CancelledError:
                raise
            except Exception as e:
                print(f"  连接异常 {type(e).__name__}: {e}，1s 后重连", flush=True)
            # 重连前固定歇一下，避免服务端持续拒绝时打成重连风暴
            if time.time() < deadline:
                await asyncio.sleep(1)
    finally:
        fh.close()

    print(f"\n########## 消息构成（共 {n_rows} 条）##########")
    for s in SYMS:
        print(f"  {s}: {dict(cnt[s])}")

    print(f"\n########## 换根时首条消息的 action（{len(rollovers)} 次）##########")
    by = defaultdict(lambda: defaultdict(int))
    gaps = defaultdict(list)
    for r in rollovers:
        by[r["sym"]][r["first_action"]] += 1
        if "gap_to_update_ms" in r:
            gaps[r["sym"]].append(r["gap_to_update_ms"])
    for s in SYMS:
        g = sorted(gaps[s])
        med = g[len(g) // 2] if g else None
        print(f"  {s}: 首条 action 分布 {dict(by[s])}"
              + (f" · snapshot->update 中位 {med}ms（n={len(g)}）" if g else ""))

    all_g = sorted(x for v in gaps.values() for x in v)
    if all_g:
        print(f"\n合计 snapshot->update 中位 {all_g[len(all_g) // 2]}ms "
              f"（n={len(all_g)}）——这就是 Hummingbot 因丢弃 snapshot 白等的时间")
    print(f"\n产物写入 {OUT}")


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--minutes", type=int, default=6)
    asyncio.run(run(ap.parse_args().minutes))


if __name__ == "__main__":
    main()