"""验证补丁是否真把那 1.1 秒拿回来了。

容器内同一进程并行跑三条路径，共用时钟与网络：
  stock   上游 BitgetPerpetualCandles（取 data[0]）
  patched PatchedBitgetPerpetualCandles（处理全部元素）
  raw     原始 aiohttp WS，按 data[-1] 判断换根（理论最快）

预期：patched ≈ raw，且比 stock 早约 1.1 秒。
同时校验补丁没有破坏数据：两条 feed 的历史 K 线应逐根相等（补丁只影响
最新一根的到达时刻与收盘价更新时机，不该改动已收盘的历史）。

    docker run --rm -v $PWD:/repo:ro -v $PWD/research/out:/out \
      -w /home/hummingbot -e PYTHONPATH=/home/hummingbot:/repo/research/live \
      --entrypoint /opt/conda/envs/hummingbot/bin/python \
      hummingbot/hummingbot:latest /repo/research/live/verify_patch.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:
    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))
                last: dict[str, int] = {}
                while not stop.is_set():
                    msg = await ws.receive()
                    if msg.type is not aiohttp.WSMsgType.TEXT:
                        break
                    if msg.data == "pong":
                        continue
                    d = json.loads(msg.data)
                    if "data" not in d or "arg" not in d or not d["data"]:
                        continue
                    sym = d["arg"]["instId"].replace("USDT", "")
                    kts = int(d["data"][-1][0])      # 关键：取末元素
                    if last.get(sym) is not None and kts > last[sym]:
                        rec.setdefault((sym, kts), {})["raw"] = \
                            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 poll_feeds(feeds: dict, key: str, rec: dict,
                     stop: asyncio.Event) -> None:
    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]:
                kts = newest * 1000 if newest < 1e12 else newest
                rec.setdefault((s, kts), {})[key] = int(time.time() * 1000)
            last[s] = newest
        await asyncio.sleep(0.01)


async def main_async(minutes: int) -> None:
    from hummingbot.data_feed.candles_feed.bitget_perpetual_candles import (
        BitgetPerpetualCandles,
    )
    from patched_candles import PatchedBitgetPerpetualCandles

    stock, patched = {}, {}
    for s in SYMS:
        stock[s] = BitgetPerpetualCandles(f"{s}-USDT", "1m", 20)
        patched[s] = PatchedBitgetPerpetualCandles(f"{s}-USDT", "1m", 20)
    for d in (stock, patched):
        for f in d.values():
            f.start()
    print(f"[补丁验证] {SYMS} · 跑 {minutes} 分钟", flush=True)

    t0 = time.time()
    while time.time() - t0 < 120:
        if all(f.ready for d in (stock, patched) for f in d.values()):
            break
        await asyncio.sleep(0.5)
    print(f"  回填完成 {time.time() - t0:.1f}s", flush=True)

    rec: dict = {}
    stop = asyncio.Event()
    tasks = [asyncio.create_task(raw_ws(rec, stop)),
             asyncio.create_task(poll_feeds(stock, "stock", rec, stop)),
             asyncio.create_task(poll_feeds(patched, "patched", rec, stop))]

    deadline = time.time() + minutes * 60
    seen = set()
    while time.time() < deadline:
        await asyncio.sleep(2)
        for k, v in rec.items():
            if k in seen or not {"stock", "patched", "raw"} <= v.keys():
                continue
            seen.add(k)
            print(f"  {k[0]}: patched 比 stock 早 "
                  f"{v['stock'] - v['patched']}ms · 距 raw "
                  f"{v['patched'] - v['raw']}ms", flush=True)
    stop.set()
    for t in tasks:
        t.cancel()
    await asyncio.gather(*tasks, return_exceptions=True)

    # 数据一致性：两条 feed 的已收盘历史必须逐根相同
    print("\n########## 补丁是否改动了已收盘 K 线 ##########")
    for s in SYMS:
        a = [list(map(float, r)) for r in list(stock[s]._candles)[:-1]]
        b = [list(map(float, r)) for r in list(patched[s]._candles)[:-1]]
        n = min(len(a), len(b))
        ta = {int(r[0]): r for r in a[-n:]}
        tb = {int(r[0]): r for r in b[-n:]}
        common = sorted(set(ta) & set(tb))
        diff = [t for t in common if ta[t][1:6] != tb[t][1:6]]
        print(f"  {s}: 共有 {len(common)} 根，OHLCV 不同 {len(diff)} 根"
              + (f"（示例 ts={diff[:3]}）" if diff else ""))
    for d in (stock, patched):
        for f in d.values():
            f.stop()

    full = [(s, k, v) for (s, k), v in rec.items()
            if {"stock", "patched", "raw"} <= v.keys()]
    f = out_dir() / "verify_patch.csv"
    with f.open("w", newline="") as fh:
        w = csv.writer(fh)
        w.writerow(["sym", "kline_ts", "raw", "patched", "stock",
                    "gain_ms", "patched_minus_raw_ms"])
        for s, k, v in sorted(full, key=lambda x: x[1]):
            w.writerow([s, k, v["raw"], v["patched"], v["stock"],
                        v["stock"] - v["patched"], v["patched"] - v["raw"]])

    print(f"\n########## 补丁收益（配对 {len(full)} 根）##########")
    print(f"{'币':<5}{'n':>5}{'早于stock中位':>14}{'距raw中位':>12}")
    for s in SYMS:
        g = [v for ss, _, v in full if ss == s]
        if not g:
            print(f"  {s}: 无样本")
            continue
        gain = sorted(v["stock"] - v["patched"] for v in g)
        dr = sorted(v["patched"] - v["raw"] for v in g)
        print(f"{s:<5}{len(g):>5}{gain[len(gain) // 2]:>14}"
              f"{dr[len(dr) // 2]:>12}")
    print(f"\n产物写入 {f}")


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


if __name__ == "__main__":
    main()