"""定位 Hummingbot 那 1030ms 究竟出在处理链路还是容器网络。

前一轮对照有两个变量同时在变：运行时（Hummingbot vs ccxt.pro）和位置
（容器内 vs 宿主机）。所以 1030ms 无法归因。原始 WS 探针已否掉「丢弃
snapshot」这个猜测——换根首条推送就是 update，Hummingbot 确实收到了。

本脚本在**容器内同一个进程**里并行跑两条路径，共用一个时钟、一条网络：
  A. Hummingbot 的 BitgetPerpetualCandles，轮询 _candles 尾部时间戳
  B. 一条原始 aiohttp WS，直接订阅 candle1m

两条路径对同一根 K 线各记一个到达时刻，差值即 Hummingbot 处理链路的净开销。
若 B 也慢，则是容器网络，与框架无关。

在容器内运行：
    docker run --rm -v $PWD:/repo:ro -v $PWD/research/out:/out \
      -w /home/hummingbot -e PYTHONPATH=/home/hummingbot \
      --entrypoint /opt/conda/envs/hummingbot/bin/python \
      hummingbot/hummingbot:latest /repo/research/live/probe_hb_vs_raw.py --minutes 20
"""
from __future__ import annotations

import argparse
import asyncio
import csv
import json
import time
from pathlib import Path

SYMS = ("BTC", "ETH", "SOL")
WSS = "wss://ws.bitget.com/v2/ws/public"


def out_dir() -> Path:
    p = Path("/out")
    return p if p.is_dir() else Path(__file__).resolve().parents[1] / "out"


async def raw_ws(rec: dict, stop: asyncio.Event) -> None:
    """B 路径：原始 WS，收到第一条带新时间戳的推送就记时刻。"""
    import aiohttp

    payload = {"op": "subscribe",
               "args": [{"instType": "USDT-FUTURES", "channel": "candle1m",
                         "instId": f"{s}USDT"} for s in SYMS]}
    while not stop.is_set():
        try:
            async with aiohttp.ClientSession() as sess, \
                    sess.ws_connect(WSS, heartbeat=20) as ws:
                await ws.send_str(json.dumps(payload))
                print("  [raw] 已订阅", flush=True)
                last: dict[str, int] = {}
                while not stop.is_set():
                    msg = await ws.receive()
                    if msg.type is not aiohttp.WSMsgType.TEXT:
                        # 关闭类消息必须跳出重连，否则 receive() 会立刻返回造成空转
                        print(f"  [raw] 非文本消息 {msg.type}，重连", flush=True)
                        break
                    if msg.data == "pong":
                        continue
                    d = json.loads(msg.data)
                    if "data" not in d or "arg" not in d:
                        continue
                    sym = d["arg"]["instId"].replace("USDT", "")
                    kts = int(d["data"][0][0])
                    if last.get(sym) is not None and kts > last[sym]:
                        rec.setdefault((sym, kts), {})["raw_ms"] = \
                            int(time.time() * 1000)
                    last[sym] = max(kts, last.get(sym, 0))
        except asyncio.CancelledError:
            raise
        except Exception as e:
            print(f"  [raw] 异常 {type(e).__name__}: {e}", flush=True)
        # 无论是正常跳出还是异常，重连前都歇一下，避免服务端持续拒绝时打成风暴
        if not stop.is_set():
            await asyncio.sleep(1)


async def hb_feed(rec: dict, stop: asyncio.Event) -> None:
    """A 路径：Hummingbot candles feed，10ms 轮询 deque 尾部。"""
    from hummingbot.data_feed.candles_feed.candles_factory import CandlesFactory
    from hummingbot.data_feed.candles_feed.data_types import CandlesConfig

    feeds = {}
    for s in SYMS:
        cfg = CandlesConfig(connector="bitget_perpetual",
                            trading_pair=f"{s}-USDT", interval="1m",
                            max_records=20)
        f = CandlesFactory.get_candle(cfg)
        f.start()
        feeds[s] = f
    print("  [hb] 已订阅", flush=True)

    t0 = time.time()
    while time.time() - t0 < 120 and not all(f.ready for f in feeds.values()):
        await asyncio.sleep(0.5)
    print(f"  [hb] 回填完成 {time.time() - t0:.1f}s", flush=True)

    last = {s: (int(f._candles[-1][0]) if len(f._candles) else None)
            for s, f in feeds.items()}
    while not stop.is_set():
        for s, f in feeds.items():
            if not len(f._candles):
                continue
            newest = int(f._candles[-1][0])
            if last[s] is not None and newest > last[s]:
                # HB 的时间戳是秒，统一成毫秒好与原始 WS 的 kline_ts 对齐
                kts = newest * 1000 if newest < 1e12 else newest
                rec.setdefault((s, kts), {})["hb_ms"] = int(time.time() * 1000)
            last[s] = newest
        await asyncio.sleep(0.01)
    for f in feeds.values():
        f.stop()


async def main_async(minutes: int, raw_only: bool = False) -> None:
    rec: dict = {}
    stop = asyncio.Event()
    tasks = [asyncio.create_task(raw_ws(rec, stop))]
    if not raw_only:
        tasks.append(asyncio.create_task(hb_feed(rec, stop)))
    where = "宿主机 raw 单跑" if raw_only else "容器内同进程对照"
    print(f"[{where}] {SYMS} · 跑 {minutes} 分钟", flush=True)

    deadline = time.time() + minutes * 60
    reported = set()
    while time.time() < deadline:
        await asyncio.sleep(2)
        for k, v in rec.items():
            if k in reported or "raw_ms" not in v or "hb_ms" not in v:
                continue
            reported.add(k)
            print(f"  {k[0]} @{k[1]}: raw->hb 延后 {v['hb_ms'] - v['raw_ms']}ms",
                  flush=True)
    stop.set()
    for t in tasks:
        t.cancel()
    await asyncio.gather(*tasks, return_exceptions=True)

    if raw_only:
        # 只落盘 raw 到达时刻，供与容器侧按 kline_ts 对齐
        f = out_dir() / "probe_raw_host.csv"
        with f.open("w", newline="") as fh:
            w = csv.writer(fh)
            w.writerow(["sym", "kline_ts", "raw_ms"])
            for (s, k), v in sorted(rec.items(), key=lambda x: x[0][1]):
                if "raw_ms" in v:
                    w.writerow([s, k, v["raw_ms"]])
        print(f"\n宿主机 raw 样本 {sum('raw_ms' in v for v in rec.values())} 根")
        print(f"产物写入 {f}")
        return

    f = out_dir() / "probe_hb_vs_raw.csv"
    both = [(s, k, v) for (s, k), v in rec.items()
            if "raw_ms" in v and "hb_ms" in v]
    with f.open("w", newline="") as fh:
        w = csv.writer(fh)
        w.writerow(["sym", "kline_ts", "raw_ms", "hb_ms", "delta_ms"])
        for s, k, v in sorted(both, key=lambda x: x[1]):
            w.writerow([s, k, v["raw_ms"], v["hb_ms"],
                        v["hb_ms"] - v["raw_ms"]])

    print(f"\n########## 同进程内 raw WS 与 Hummingbot 的到达差 ##########")
    print(f"（正数 = Hummingbot 更慢；配对 {len(both)} 根）")
    for s in SYMS:
        d = sorted(v["hb_ms"] - v["raw_ms"] for ss, _, v in both if ss == s)
        if not d:
            print(f"  {s}: 无配对样本")
            continue
        print(f"  {s}: n={len(d)} 中位 {d[len(d) // 2]}ms "
              f"P90 {d[int(len(d) * .9)]}ms 最小 {d[0]}ms 最大 {d[-1]}ms")
    print(f"\n判读：中位接近 0 说明 1030ms 来自容器网络或宿主机对照本身；"
          f"中位接近 1000ms 说明是 Hummingbot 处理链路的净开销。")
    print(f"产物写入 {f}")


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--minutes", type=int, default=20)
    ap.add_argument("--raw-only", action="store_true",
                    help="不加载 hummingbot，只跑原始 WS（供宿主机对照）")
    a = ap.parse_args()
    asyncio.run(main_async(a.minutes, a.raw_only))


if __name__ == "__main__":
    main()