crypto_okx_hub/lib/cash_sample_lib.py

"""中控余额采样(曲线主指标=现金,不含权利金/浮盈)."""
from __future__ import annotations

import os
import threading
import time
from typing import Any

from lib.db_lib import get_db

_lock = threading.Lock()
_last_sample_at: dict[Any, float] = {}


def cash_sample_interval_sec() -> float:
    try:
        return max(30.0, float(os.getenv("HUB_CASH_SAMPLE_SEC") or 60))
    except (TypeError, ValueError):
        return 60.0


def cash_sample_keep_days() -> int:
    try:
        return max(7, int(float(os.getenv("HUB_CASH_SAMPLE_KEEP_DAYS") or 90)))
    except (TypeError, ValueError):
        return 90


def _safe_float(v: Any) -> float | None:
    if v is None or v == "":
        return None
    try:
        return float(v)
    except (TypeError, ValueError):
        return None


def cash_from_overview_row(row: dict[str, Any]) -> float | None:
    snap = row.get("snapshot") if isinstance(row.get("snapshot"), dict) else {}
    premium = snap.get("premium") if isinstance(snap.get("premium"), dict) else {}
    cash = _safe_float(premium.get("cash_usdt"))
    if cash is not None:
        return round(cash, 4)
    funds = snap.get("funds") if isinstance(snap.get("funds"), dict) else {}
    parts = [_safe_float(funds.get("trading_usdt")), _safe_float(funds.get("funding_usdt"))]
    if all(x is None for x in parts):
        parts = [_safe_float(funds.get("trading_usdc")), _safe_float(funds.get("funding_usdc"))]
    known = [x for x in parts if x is not None]
    return round(sum(known), 4) if known else None


def record_cash_samples_from_overview(rows: list[dict[str, Any]]) -> int:
    """从总览拉取结果按间隔写入余额样本.返回写入条数."""
    if not rows:
        return 0
    interval = cash_sample_interval_sec()
    now = time.time()
    ts_ms = int(now * 1000)
    to_write: list[tuple[int, int, float]] = []
    with _lock:
        for row in rows:
            if not row.get("online"):
                continue
            iid = row.get("id")
            if iid is None:
                continue
            cash = cash_from_overview_row(row)
            if cash is None:
                continue
            last = float(_last_sample_at.get(iid) or 0)
            if now - last < interval:
                continue
            _last_sample_at[iid] = now
            to_write.append((int(iid), ts_ms, float(cash)))
    if not to_write:
        return 0
    conn = get_db()
    try:
        conn.executemany(
            "INSERT OR IGNORE INTO hub_cash_samples(instance_id, ts_ms, cash_usdt) VALUES (?,?,?)",
            to_write,
        )
        # 偶尔清理过期
        if int(now) % 50 < len(to_write):
            keep_ms = int(now * 1000) - cash_sample_keep_days() * 86400 * 1000
            conn.execute("DELETE FROM hub_cash_samples WHERE ts_ms < ?", (keep_ms,))
        conn.commit()
    finally:
        conn.close()
    return len(to_write)


def query_cash_series(
    *,
    range_key: str = "7d",
    instance_ids: list[int] | None = None,
) -> dict[str, Any]:
    """返回 {series: [{id,name,points:[{t,v}]}], total:[{t,v}]}."""
    from lib.instances_lib import list_instances

    rk = (range_key or "7d").strip().lower()
    now_ms = int(time.time() * 1000)
    if rk in ("today", "1d"):
        # 今日 0 点起(本地)
        import datetime as dt

        start = dt.datetime.now().replace(hour=0, minute=0, second=0, microsecond=0)
        start_ms = int(start.timestamp() * 1000)
    elif rk in ("30d", "30"):
        start_ms = now_ms - 30 * 86400 * 1000
    elif rk in ("all", "0"):
        start_ms = 0
    else:
        start_ms = now_ms - 7 * 86400 * 1000

    instances = list_instances(enabled_only=True)
    if instance_ids is not None:
        want = {int(x) for x in instance_ids}
        instances = [i for i in instances if int(i["id"]) in want]
    id_name = {int(i["id"]): i.get("name") or str(i["id"]) for i in instances}
    ids = list(id_name.keys())
    if not ids:
        return {"series": [], "total": [], "range": rk, "start_ms": start_ms, "end_ms": now_ms}

    placeholders = ",".join("?" * len(ids))
    conn = get_db()
    try:
        rows = conn.execute(
            f"""
            SELECT instance_id, ts_ms, cash_usdt
            FROM hub_cash_samples
            WHERE instance_id IN ({placeholders})
              AND ts_ms >= ?
            ORDER BY ts_ms ASC
            """,
            (*ids, start_ms),
        ).fetchall()
    finally:
        conn.close()

    by_id: dict[int, list[dict[str, float]]] = {i: [] for i in ids}
    for r in rows:
        iid = int(r["instance_id"])
        t = int(r["ts_ms"])
        v = float(r["cash_usdt"])
        by_id.setdefault(iid, []).append({"t": t, "v": v})

    # 合计:按 1 分钟桶,取该桶内各机最后样本求和,并前向填充
    last_in_bucket: dict[int, dict[int, float]] = {}
    for iid, pts in by_id.items():
        for p in pts:
            bt = (int(p["t"]) // 60000) * 60000
            last_in_bucket.setdefault(bt, {})[iid] = float(p["v"])
    running: dict[int, float] = {}
    total_pts: list[dict[str, float]] = []
    for bt in sorted(last_in_bucket.keys()):
        running.update(last_in_bucket[bt])
        if running:
            total_pts.append({"t": bt, "v": round(sum(running.values()), 4)})

    series = [
        {
            "id": iid,
            "name": id_name.get(iid, str(iid)),
            "points": by_id.get(iid) or [],
        }
        for iid in ids
    ]
    return {
        "series": series,
        "total": total_pts,
        "range": rk,
        "start_ms": start_ms,
        "end_ms": now_ms,
    }