crypto_okx/lib/sim/hooks.py

"""运行时挂钩: 将 app / options cfg 在 sim 模式下切到本地撮合."""

from __future__ import annotations

from typing import Any, Callable, Optional

from lib.sim.broker_lib import SimBroker
from lib.sim.mode_lib import is_sim_mode
from lib.sim.pricing_lib import sim_fee_rate


_GET_DB: Optional[Callable] = None
_APP_MODULE: Any = None


def set_get_db(get_db: Callable) -> None:
    global _GET_DB
    _GET_DB = get_db


def get_db_fn() -> Callable:
    if _GET_DB is None:
        raise RuntimeError("sim get_db 未安装")
    return _GET_DB


def broker() -> SimBroker:
    return SimBroker(get_db_fn())


def apply_sim_hooks(app_module: Any) -> None:
    """包装 app 上的永续下单/余额/就绪检查; 幂等."""
    global _APP_MODULE
    _APP_MODULE = app_module
    set_get_db(app_module.get_db)

    if getattr(app_module, "_sim_hooks_applied", False):
        return

    _orig_ensure = app_module.ensure_okx_live_ready
    _orig_capitals = app_module.get_exchange_capitals
    _orig_avail = app_module.get_available_trading_usdt
    _orig_place = app_module.place_exchange_order
    _orig_close = app_module.close_exchange_order
    _orig_live_contracts = app_module.get_live_position_contracts

    def ensure_okx_live_ready():
        if is_sim_mode(app_module.get_db):
            return True, "sim"
        return _orig_ensure()

    def get_exchange_capitals(force=False):
        if is_sim_mode(app_module.get_db):
            w = broker().balances_header()
            return float(w["funding_usdt"]), float(w["trading_usdt"])
        return _orig_capitals(force=force)

    def get_available_trading_usdt():
        if is_sim_mode(app_module.get_db):
            return float(broker().balances_header()["trading_usdt"])
        return _orig_avail()

    def place_exchange_order(
        exchange_symbol, direction, amount, leverage, stop_loss=None, take_profit=None
    ):
        if is_sim_mode(app_module.get_db):
            ex = getattr(app_module, "exchange", None)
            if ex is None:
                raise RuntimeError("sim: exchange 未就绪(公开行情)")
            ensure = getattr(app_module, "ensure_markets_loaded", None)
            if callable(ensure):
                try:
                    ensure()
                except Exception:
                    pass
            return broker().place_perp_market(
                ex,
                symbol=exchange_symbol,
                direction=direction,
                contracts=float(amount),
                leverage=int(leverage or 1),
                fee_rate=sim_fee_rate(),
                stop_loss=stop_loss,
                take_profit=take_profit,
            )
        return _orig_place(
            exchange_symbol, direction, amount, leverage, stop_loss=stop_loss, take_profit=take_profit
        )

    def close_exchange_order(order_row):
        if is_sim_mode(app_module.get_db):
            ex = getattr(app_module, "exchange", None)
            if ex is None:
                raise RuntimeError("sim: exchange 未就绪")
            ensure = getattr(app_module, "ensure_markets_loaded", None)
            if callable(ensure):
                try:
                    ensure()
                except Exception:
                    pass
            normalize = getattr(app_module, "normalize_okx_symbol", None) or getattr(
                app_module, "normalize_exchange_symbol", None
            )
            try:
                symbol = order_row["exchange_symbol"] or None
            except Exception:
                symbol = None
            if not symbol:
                try:
                    symbol = order_row["symbol"]
                except Exception:
                    symbol = None
            if callable(normalize):
                symbol = normalize(symbol)
            direction = order_row["direction"]
            db_amt = float(order_row["order_amount"] or 0)
            live = broker().get_perp_contracts(symbol, direction)
            amt = live if live and live > 0 else db_amt
            return broker().close_perp_market(
                ex, symbol=symbol, direction=direction, contracts=amt, fee_rate=sim_fee_rate()
            )
        return _orig_close(order_row)

    def get_live_position_contracts(exchange_symbol, direction):
        if is_sim_mode(app_module.get_db):
            normalize = getattr(app_module, "normalize_okx_symbol", None)
            sym = exchange_symbol
            if callable(normalize):
                sym = normalize(exchange_symbol or "")
            return broker().get_perp_contracts(sym, direction)
        return _orig_live_contracts(exchange_symbol, direction)

    app_module.ensure_okx_live_ready = ensure_okx_live_ready
    app_module.get_exchange_capitals = get_exchange_capitals
    app_module.get_available_trading_usdt = get_available_trading_usdt
    app_module.place_exchange_order = place_exchange_order
    app_module.close_exchange_order = close_exchange_order
    app_module.get_live_position_contracts = get_live_position_contracts
    app_module._sim_hooks_applied = True

    _patch_okx_options_lib(app_module)


def _patch_okx_options_lib(app_module: Any) -> None:
    """期权余额 / 成交等待: 对 sim-* 订单与 sim 模式短路."""
    import lib.exchange.okx_options_lib as opt_lib

    if getattr(opt_lib, "_sim_hooks_applied", False):
        return

    _orig_header = opt_lib.options_header_balances
    _orig_wait = opt_lib.wait_option_order_full_fill
    _orig_fetch_order = opt_lib.fetch_option_order
    _orig_fetch_pos = opt_lib.fetch_option_positions
    _orig_ready = opt_lib.options_api_ready
    _orig_fetch_bal = opt_lib.fetch_options_balances

    def options_header_balances(ex, *, force: bool = False):
        try:
            if _GET_DB is not None and is_sim_mode(_GET_DB):
                w = broker().balances_header()
                return (
                    round(float(w["trading_usdc"]), 2),
                    round(float(w["funding_usdc"]), 2),
                    round(float(w["funding_usdt"]), 2),
                    round(float(w["trading_usdt"]), 2),
                )
        except Exception:
            pass
        return _orig_header(ex, force=force)

    def fetch_options_balances(ex, *, force: bool = False):
        try:
            if _GET_DB is not None and is_sim_mode(_GET_DB):
                w = broker().balances_header()
                return {
                    "trading_usdc": float(w["trading_usdc"]),
                    "funding_usdc": float(w["funding_usdc"]),
                    "funding_usdt": float(w["funding_usdt"]),
                    "trading_usdt": float(w["trading_usdt"]),
                }
        except Exception:
            pass
        return _orig_fetch_bal(ex, force=force)

    def options_api_ready(ex):
        try:
            if _GET_DB is not None and is_sim_mode(_GET_DB):
                return True, "sim"
        except Exception:
            pass
        return _orig_ready(ex)

    def fetch_option_order(ex, *, inst_id: str, ord_id: str):
        oid = str(ord_id or "")
        if oid.startswith("sim-"):
            info = broker().get_option_order(oid)
            if info:
                return info
            return {"ok": False, "msg": "sim order not found"}
        return _orig_fetch_order(ex, inst_id=inst_id, ord_id=ord_id)

    def wait_option_order_full_fill(
        ex,
        *,
        inst_id: str,
        ord_id: str,
        need_sheets: int,
        timeout_sec: float = 12.0,
        poll_sec: float = 0.35,
        cancel_on_timeout: bool = True,
    ):
        oid = str(ord_id or "")
        if oid.startswith("sim-"):
            info = broker().get_option_order(oid)
            if not info:
                return {"ok": False, "msg": "sim order not found", "filled_sheets": 0}
            return {
                "ok": True,
                "filled_sheets": int(round(float(info.get("acc_fill_sz") or need_sheets))),
                "avg_px": info.get("avg_px"),
                "state": "filled",
                "order": info,
            }
        return _orig_wait(
            ex,
            inst_id=inst_id,
            ord_id=ord_id,
            need_sheets=need_sheets,
            timeout_sec=timeout_sec,
            poll_sec=poll_sec,
            cancel_on_timeout=cancel_on_timeout,
        )

    def fetch_option_positions(ex):
        try:
            if _GET_DB is not None and is_sim_mode(_GET_DB):
                return broker().option_positions_okx_rows()
        except Exception:
            pass
        return _orig_fetch_pos(ex)

    opt_lib.options_header_balances = options_header_balances
    opt_lib.fetch_options_balances = fetch_options_balances
    opt_lib.options_api_ready = options_api_ready
    opt_lib.fetch_option_order = fetch_option_order
    opt_lib.wait_option_order_full_fill = wait_option_order_full_fill
    opt_lib.fetch_option_positions = fetch_option_positions
    opt_lib._sim_hooks_applied = True


def wrap_option_place_fns(get_db: Callable, live_place_limit, live_place_market):
    """返回按模式分流的 place_option_limit/market."""

    def place_option_limit_order(ex, **kwargs):
        if is_sim_mode(get_db):
            side = (kwargs.get("side") or "").lower()
            b = broker()
            if side == "sell" or kwargs.get("reduce_only"):
                return b.sell_option_close(ex, **kwargs)
            return b.place_option_buy(ex, **kwargs)
        return live_place_limit(ex, **kwargs)

    def place_option_market_order(ex, **kwargs):
        if is_sim_mode(get_db):
            side = (kwargs.get("side") or "").lower()
            b = broker()
            if side == "sell" or kwargs.get("reduce_only"):
                return b.sell_option_close(ex, **kwargs)
            return b.place_option_buy(ex, **kwargs)
        return live_place_market(ex, **kwargs)

    return place_option_limit_order, place_option_market_order


def patch_options_cfg(cfg: dict[str, Any]) -> dict[str, Any]:
    """就地替换 options/hedge cfg 中的下单函数为模式感知包装."""
    get_db = cfg.get("get_db")
    if not callable(get_db):
        return cfg
    set_get_db(get_db)
    live_limit = cfg.get("place_option_limit_order")
    live_market = cfg.get("place_option_market_order")
    if callable(live_limit) and callable(live_market):
        wrapped_l, wrapped_m = wrap_option_place_fns(get_db, live_limit, live_market)
        cfg["place_option_limit_order"] = wrapped_l
        cfg["place_option_market_order"] = wrapped_m
    elif callable(live_limit):
        wrapped_l, _ = wrap_option_place_fns(
            get_db,
            live_limit,
            live_limit,
        )
        cfg["place_option_limit_order"] = wrapped_l

    live_cancel = cfg.get("cancel_option_order")

    def cancel_option_order(ex, **kwargs):
        oid = str(kwargs.get("ord_id") or "")
        if oid.startswith("sim-") or (callable(get_db) and is_sim_mode(get_db)):
            return {"ok": True, "msg": "sim cancel noop", "sim": True}
        if callable(live_cancel):
            return live_cancel(ex, **kwargs)
        return {"ok": False, "msg": "cancel unavailable"}

    if "cancel_option_order" in cfg:
        cfg["cancel_option_order"] = cancel_option_order

    live_pending = cfg.get("fetch_option_pending_orders")

    def fetch_option_pending_orders(ex, **kwargs):
        if is_sim_mode(get_db):
            return []
        if callable(live_pending):
            return live_pending(ex, **kwargs)
        return []

    if "fetch_option_pending_orders" in cfg:
        cfg["fetch_option_pending_orders"] = fetch_option_pending_orders

    return cfg