#!/usr/bin/env python3
"""Decode 178-byte FundFlowNats messages from the stock fund-flow service."""

import argparse
import asyncio
import datetime as dt

import numpy as np
from nats.aio.client import Client as NATS


FUNDFLOW = np.dtype([
    ("symbol", "S32"), ("trade_date", "i4"), ("market", "u1"), ("bucket_ms", "i8"),
    ("version", "u8"), ("status", "u1"),
    ("main_inflow_x1e4", "i8"), ("main_outflow_x1e4", "i8"), ("main_net_x1e4", "i8"),
    ("ultra_inflow_x1e4", "i8"), ("ultra_outflow_x1e4", "i8"), ("ultra_net_x1e4", "i8"),
    ("large_inflow_x1e4", "i8"), ("large_outflow_x1e4", "i8"), ("large_net_x1e4", "i8"),
    ("medium_inflow_x1e4", "i8"), ("medium_outflow_x1e4", "i8"), ("medium_net_x1e4", "i8"),
    ("small_inflow_x1e4", "i8"), ("small_outflow_x1e4", "i8"), ("small_net_x1e4", "i8"),
    ("event_count", "u4"),
])
assert FUNDFLOW.itemsize == 178
STATUS = {0: "临时", 1: "修正", 2: "最终"}
MARKET = {1: "SH", 2: "SZ"}


def bucket_time(bucket_ms: int) -> str:
    seconds = bucket_ms // 1000
    return f"{seconds // 3600:02d}:{(seconds % 3600) // 60:02d}:{seconds % 60:02d}"


def yuan(value: int) -> float:
    return value / 10000.0


async def run(args: argparse.Namespace) -> None:
    nc = NATS()
    await nc.connect(
        servers=[args.url], user=args.user, password=args.password,
        allow_reconnect=True, max_reconnect_attempts=-1, reconnect_time_wait=2,
    )

    async def on_message(msg):
        if len(msg.data) != FUNDFLOW.itemsize:
            print(f"[-] {msg.subject}: expected 178 bytes, got {len(msg.data)}")
            return
        rec = np.frombuffer(msg.data, dtype=FUNDFLOW)[0]
        symbol = rec["symbol"].decode("ascii", errors="replace").rstrip("\x00")
        print(
            f"[{dt.datetime.now():%H:%M:%S.%f}] {msg.subject} {symbol} "
            f"{MARKET.get(int(rec['market']), rec['market'])} {bucket_time(int(rec['bucket_ms']))} "
            f"status={STATUS.get(int(rec['status']), rec['status'])} version={rec['version']} "
            f"main_net={yuan(int(rec['main_net_x1e4'])):+.2f} yuan events={rec['event_count']}"
        )

    await nc.subscribe(args.subject, cb=on_message, pending_msgs_limit=1_000_000, pending_bytes_limit=256 * 1024 * 1024)
    print(f"[+] subscribed: {args.subject}")
    try:
        while True:
            await asyncio.sleep(1)
    finally:
        await nc.drain()


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--url", default="nats://quotex.base32.cn:4222")
    parser.add_argument("--user", required=True)
    parser.add_argument("--password", required=True)
    parser.add_argument("--subject", default="stk.flow.1s.600519")
    return parser.parse_args()


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