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

    _orig_xfer = getattr(app_module, "execute_transfer_usdt", None)
    if callable(_orig_xfer):

        def execute_transfer_usdt(amount, from_account, to_account):
            if is_sim_mode(app_module.get_db):
                from lib.sim.wallets_lib import SimWallets, normalize_sim_account

                try:
                    result = SimWallets(app_module.get_db).transfer(
                        ccy="USDT",
                        amount=float(amount),
                        from_account=normalize_sim_account(from_account),
                        to_account=normalize_sim_account(to_account),
                    )
                except Exception as e:
                    return False, str(e), None
                if not result.get("ok"):
                    return False, result.get("detail") or result.get("msg") or "划转失败", None
                try:
                    from lib.instance.instance_live_push_lib import notify_instance_balance_changed

                    notify_instance_balance_changed()
                except Exception:
                    pass
                return True, "sim 划转成功", result
            return _orig_xfer(amount, from_account, to_account)

        app_module.execute_transfer_usdt = execute_transfer_usdt

    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
    _orig_transfer = opt_lib.transfer_ccy
    _orig_swap = opt_lib.spot_market_swap_usdt_usdc

    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, scope: str = "main", sub_acct: str = ""
    ):
        try:
            if _GET_DB is not None and is_sim_mode(_GET_DB):
                w = broker().balances_header()
                fu = float(w["funding_usdt"])
                fc = float(w["funding_usdc"])
                tu = float(w["trading_usdt"])
                tc = float(w["trading_usdc"])
                return {
                    "scope": "main",
                    "funding_usdt": fu,
                    "funding_usdc": fc,
                    "trading_usdt": tu,
                    "trading_usdc": tc,
                    "funding_usdt_avail": fu,
                    "funding_usdc_avail": fc,
                    "trading_usdt_avail": tu,
                    "trading_usdc_avail": tc,
                }
        except Exception:
            pass
        return _orig_fetch_bal(ex, force=force, scope=scope, sub_acct=sub_acct)

    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 transfer_ccy(ex, ccy, amount, from_acct, to_acct):
        try:
            if _GET_DB is not None and is_sim_mode(_GET_DB):
                from lib.sim.wallets_lib import SimWallets, normalize_sim_account

                result = SimWallets(_GET_DB).transfer(
                    ccy=str(ccy or "USDC"),
                    amount=float(amount),
                    from_account=normalize_sim_account(from_acct),
                    to_account=normalize_sim_account(to_acct),
                )
                if not result.get("ok"):
                    return {
                        "ok": False,
                        "msg": result.get("detail") or result.get("msg") or "划转失败",
                    }
                try:
                    from lib.instance.instance_live_push_lib import notify_instance_balance_changed

                    notify_instance_balance_changed()
                except Exception:
                    pass
                return {"ok": True, "msg": "sim 划转成功", "sim": True, **result}
        except Exception as e:
            return {"ok": False, "msg": str(e)}
        return _orig_transfer(ex, ccy, amount, from_acct, to_acct)

    def spot_market_swap_usdt_usdc(ex, *, direction: str = "usdt_to_usdc", amount: float = 0):
        try:
            if _GET_DB is not None and is_sim_mode(_GET_DB):
                from lib.sim.wallets_lib import SimWallets

                d = (direction or "usdt_to_usdc").strip().lower()
                if d == "usdc_to_usdt":
                    from_ccy, to_ccy = "USDC", "USDT"
                else:
                    from_ccy, to_ccy = "USDT", "USDC"
                result = SimWallets(_GET_DB).convert(
                    from_ccy=from_ccy,
                    to_ccy=to_ccy,
                    amount=float(amount),
                    account="funding",
                )
                if not result.get("ok"):
                    result = SimWallets(_GET_DB).convert(
                        from_ccy=from_ccy,
                        to_ccy=to_ccy,
                        amount=float(amount),
                        account="trading",
                    )
                if not result.get("ok"):
                    return {
                        "ok": False,
                        "msg": result.get("detail") or result.get("msg") or "兑换失败",
                    }
                try:
                    from lib.instance.instance_live_push_lib import notify_instance_balance_changed

                    notify_instance_balance_changed()
                except Exception:
                    pass
                return {"ok": True, "msg": "sim 兑换成功(1:1)", "sim": True, **result}
        except Exception as e:
            return {"ok": False, "msg": str(e)}
        return _orig_swap(ex, direction=direction, amount=amount)

    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.transfer_ccy = transfer_ccy
    opt_lib.spot_market_swap_usdt_usdc = spot_market_swap_usdt_usdc
    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)

    # 与 apply_sim_hooks 对齐: cfg 里可能仍是 import 时的旧引用
    import lib.exchange.okx_options_lib as opt_lib

    for key in (
        "transfer_ccy",
        "spot_market_swap_usdt_usdc",
        "options_api_ready",
        "fetch_options_balances",
        "fetch_option_positions",
        "fetch_option_order",
        "wait_option_order_full_fill",
    ):
        if key in cfg and hasattr(opt_lib, key):
            cfg[key] = getattr(opt_lib, key)

    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