crypto_okx/lib/sim/broker_lib.py

"""模拟撮合: 用公开行情 bid/ask 成交, 结算到本地钱包."""

from __future__ import annotations

import uuid
from datetime import datetime
from typing import Any, Callable

from lib.sim.pricing_lib import option_fill, perp_fill, sim_fee_rate
from lib.sim.wallets_lib import InsufficientFunds, SimWallets


def _now() -> str:
    return datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S")


def _ticker_bid_ask(exchange: Any, symbol: str) -> tuple[float, float]:
    t = exchange.fetch_ticker(symbol)
    last = t.get("last")
    bid = t.get("bid")
    ask = t.get("ask")
    if bid is None or float(bid) <= 0:
        bid = last
    if ask is None or float(ask) <= 0:
        ask = last
    if bid is None or ask is None:
        raise RuntimeError(f"无法获取 {symbol} 行情 bid/ask")
    return float(bid), float(ask)


def _option_bid_ask(exchange_options: Any, inst_id: str) -> tuple[float, float, float]:
    """返回 bid, ask, ct_mult."""
    ct_mult = 0.01
    bid = ask = None
    try:
        from lib.exchange.okx_options_lib import quote_option_contract

        q = quote_option_contract(exchange_options, inst_id)
        if q.get("ok"):
            bid = q.get("bid") or q.get("mark")
            ask = q.get("ask") or q.get("book_ask") or q.get("ref_ask") or q.get("mark")
            ct_mult = float(q.get("ct_mult") or 0.01)
    except Exception:
        pass
    if bid is None or ask is None:
        rows = exchange_options.public_get_market_ticker({"instId": inst_id}).get("data") or []
        if not rows:
            raise RuntimeError(f"无法获取期权行情 {inst_id}")
        t = rows[0]
        bid = bid or t.get("bidPx") or t.get("markPx") or t.get("last")
        ask = ask or t.get("askPx") or t.get("markPx") or t.get("last")
    if bid is None or ask is None:
        raise RuntimeError(f"期权 {inst_id} 缺少 bid/ask")
    return float(bid), float(ask), float(ct_mult)


def _contract_size(exchange: Any, symbol: str) -> float:
    try:
        if hasattr(exchange, "market"):
            m = exchange.market(symbol)
            return float(m.get("contractSize") or 1)
    except Exception:
        pass
    return 1.0


class SimBroker:
    def __init__(self, get_db: Callable) -> None:
        self.get_db = get_db
        self.wallets = SimWallets(get_db)

    def balances_header(self) -> dict[str, float]:
        return self.wallets.view()

    def list_perp_positions(self) -> list[dict[str, Any]]:
        conn = self.get_db()
        try:
            rows = conn.execute(
                "SELECT * FROM sim_perp_positions WHERE contracts > 1e-12 ORDER BY id"
            ).fetchall()
            out = []
            for r in rows:
                out.append(
                    {
                        "symbol": r["symbol"],
                        "direction": r["direction"],
                        "contracts": float(r["contracts"]),
                        "entry_px": float(r["entry_px"]),
                        "leverage": int(r["leverage"] or 1),
                        "margin_usdt": float(r["margin_usdt"] or 0),
                        "contract_size": float(r["contract_size"] or 1),
                    }
                )
            return out
        finally:
            conn.close()

    def list_option_positions(self) -> list[dict[str, Any]]:
        conn = self.get_db()
        try:
            rows = conn.execute(
                "SELECT * FROM sim_option_positions WHERE sheets > 1e-12 ORDER BY id"
            ).fetchall()
            out = []
            for r in rows:
                out.append(
                    {
                        "inst_id": r["inst_id"],
                        "side": r["side"],
                        "sheets": float(r["sheets"]),
                        "entry_px": float(r["entry_px"]),
                        "ct_mult": float(r["ct_mult"] or 0.01),
                        "premium_paid_usdc": float(r["premium_paid_usdc"] or 0),
                    }
                )
            return out
        finally:
            conn.close()

    def get_perp_contracts(self, symbol: str, direction: str) -> float:
        conn = self.get_db()
        try:
            row = conn.execute(
                """
                SELECT contracts FROM sim_perp_positions
                WHERE symbol=? AND direction=? AND contracts > 1e-12
                """,
                (symbol, direction),
            ).fetchone()
            if not row:
                return 0.0
            return float(row["contracts"] if hasattr(row, "keys") else row[0])
        finally:
            conn.close()

    def place_perp_market(
        self,
        exchange: Any,
        *,
        symbol: str,
        direction: str,
        contracts: float,
        leverage: int,
        fee_rate: float | None = None,
        stop_loss: Any = None,
        take_profit: Any = None,
    ) -> dict[str, Any]:
        _ = (stop_loss, take_profit)  # sim: ignore attachAlgoOrds / tpsl
        side = (direction or "").lower().strip()
        if side not in ("long", "short"):
            raise ValueError("direction 须为 long 或 short")
        qty_c = float(contracts)
        if qty_c <= 0:
            raise ValueError("张数须大于 0")
        lev = max(1, int(leverage or 1))
        fr = sim_fee_rate(fee_rate)
        bid, ask = _ticker_bid_ask(exchange, symbol)
        ct_sz = _contract_size(exchange, symbol)
        qty = qty_c * ct_sz
        pr = perp_fill(side=side, action="open", bid=bid, ask=ask, qty=qty, fee_rate=fr)
        margin = pr.notional / lev
        need = margin + pr.fee
        try:
            self.wallets.debit_trading(
                "USDT",
                need,
                kind="perp_open",
                note=f"open {side} {symbol} {qty_c}@{pr.fill_px:.4f}",
            )
        except InsufficientFunds as e:
            raise RuntimeError(str(e)) from e

        conn = self.get_db()
        try:
            existing = conn.execute(
                "SELECT * FROM sim_perp_positions WHERE symbol=? AND direction=?",
                (symbol, side),
            ).fetchone()
            if existing and float(existing["contracts"] or 0) > 1e-12:
                old_c = float(existing["contracts"])
                old_px = float(existing["entry_px"])
                old_m = float(existing["margin_usdt"] or 0)
                new_c = old_c + qty_c
                entry = (old_px * old_c + pr.fill_px * qty_c) / new_c
                conn.execute(
                    """
                    UPDATE sim_perp_positions
                    SET contracts=?, entry_px=?, leverage=?, margin_usdt=?, contract_size=?, updated_at=?
                    WHERE symbol=? AND direction=?
                    """,
                    (new_c, entry, lev, old_m + margin, ct_sz, _now(), symbol, side),
                )
            else:
                conn.execute(
                    """
                    INSERT INTO sim_perp_positions(
                        symbol, direction, contracts, entry_px, leverage, margin_usdt, contract_size, updated_at
                    ) VALUES (?,?,?,?,?,?,?,?)
                    ON CONFLICT(symbol, direction) DO UPDATE SET
                        contracts=excluded.contracts,
                        entry_px=excluded.entry_px,
                        leverage=excluded.leverage,
                        margin_usdt=excluded.margin_usdt,
                        contract_size=excluded.contract_size,
                        updated_at=excluded.updated_at
                    """,
                    (symbol, side, qty_c, pr.fill_px, lev, margin, ct_sz, _now()),
                )
            conn.commit()
        finally:
            conn.close()

        oid = f"sim-perp-{uuid.uuid4().hex[:16]}"
        return {
            "id": oid,
            "symbol": symbol,
            "side": "buy" if side == "long" else "sell",
            "amount": qty_c,
            "average": pr.fill_px,
            "status": "closed",
            "info": {"sim": True, "fee": pr.fee, "margin": margin, "fill": pr.to_dict()},
            "tpsl_attached": False,
        }

    def close_perp_market(
        self,
        exchange: Any,
        *,
        symbol: str,
        direction: str,
        contracts: float | None = None,
        fee_rate: float | None = None,
    ) -> dict[str, Any]:
        side = (direction or "").lower().strip()
        if side not in ("long", "short"):
            raise ValueError("direction 须为 long 或 short")
        conn = self.get_db()
        try:
            row = conn.execute(
                "SELECT * FROM sim_perp_positions WHERE symbol=? AND direction=?",
                (symbol, side),
            ).fetchone()
            if not row or float(row["contracts"] or 0) <= 1e-12:
                raise ValueError("模拟永续无对应持仓")
            pos_c = float(row["contracts"])
            entry = float(row["entry_px"])
            margin_all = float(row["margin_usdt"] or 0)
            ct_sz = float(row["contract_size"] or 1)
            close_c = pos_c if contracts is None else min(pos_c, float(contracts))
            if close_c <= 0:
                raise ValueError("平仓张数无效")
            fr = sim_fee_rate(fee_rate)
            bid, ask = _ticker_bid_ask(exchange, symbol)
            qty = close_c * ct_sz
            pr = perp_fill(side=side, action="close", bid=bid, ask=ask, qty=qty, fee_rate=fr)
            if side == "long":
                pnl = (pr.fill_px - entry) * qty
            else:
                pnl = (entry - pr.fill_px) * qty
            release = margin_all * (close_c / pos_c)
            credit = release + pnl - pr.fee
            remain = pos_c - close_c
            if remain <= 1e-12:
                conn.execute(
                    "DELETE FROM sim_perp_positions WHERE symbol=? AND direction=?",
                    (symbol, side),
                )
            else:
                conn.execute(
                    """
                    UPDATE sim_perp_positions
                    SET contracts=?, margin_usdt=?, updated_at=?
                    WHERE symbol=? AND direction=?
                    """,
                    (remain, margin_all - release, _now(), symbol, side),
                )
            conn.commit()
        finally:
            conn.close()

        if credit >= 0:
            self.wallets.credit_trading(
                "USDT",
                credit,
                kind="perp_close",
                note=f"close {side} {symbol} pnl={pnl:.4f}",
            )
        else:
            self.wallets.debit_trading(
                "USDT",
                abs(credit),
                kind="perp_close",
                note=f"close {side} {symbol} pnl={pnl:.4f}",
            )

        oid = f"sim-perp-close-{uuid.uuid4().hex[:16]}"
        return {
            "id": oid,
            "symbol": symbol,
            "side": "sell" if side == "long" else "buy",
            "amount": close_c,
            "average": pr.fill_px,
            "status": "closed",
            "info": {
                "sim": True,
                "fee": pr.fee,
                "pnl": pnl,
                "released_margin": release,
                "fill": pr.to_dict(),
            },
            "tpsl_attached": False,
        }

    def _store_option_order(
        self,
        *,
        ord_id: str,
        inst_id: str,
        side: str,
        sheets: float,
        avg_px: float,
    ) -> None:
        conn = self.get_db()
        try:
            conn.execute(
                """
                INSERT OR REPLACE INTO sim_option_orders(
                    ord_id, inst_id, side, sheets, avg_px, state, acc_fill_sz, created_at
                ) VALUES (?,?,?,?,?,'filled',?,?)
                """,
                (ord_id, inst_id, side, float(sheets), float(avg_px), float(sheets), _now()),
            )
            conn.commit()
        finally:
            conn.close()

    def get_option_order(self, ord_id: str) -> dict[str, Any] | None:
        conn = self.get_db()
        try:
            row = conn.execute(
                "SELECT * FROM sim_option_orders WHERE ord_id=?",
                (ord_id,),
            ).fetchone()
            if not row:
                return None
            return {
                "ok": True,
                "ord_id": row["ord_id"],
                "inst_id": row["inst_id"],
                "side": row["side"],
                "state": row["state"],
                "acc_fill_sz": float(row["acc_fill_sz"]),
                "avg_px": float(row["avg_px"]),
                "sim": True,
            }
        finally:
            conn.close()

    def place_option_buy(
        self,
        exchange_options: Any,
        *,
        inst_id: str,
        sheets: int,
        price: float | None = None,
        fee_rate: float | None = None,
        **_kwargs: Any,
    ) -> dict[str, Any]:
        n = int(sheets)
        if n < 1:
            return {"ok": False, "msg": "张数至少为 1"}
        fr = sim_fee_rate(fee_rate)
        bid, ask, ct_mult = _option_bid_ask(exchange_options, inst_id)
        if price is not None and float(price) > 0:
            # 限价: 用 min(limit, ask) 作为基准卖一近似, 仍走 option_fill 滑点
            ask = min(float(ask), float(price)) if float(price) > 0 else float(ask)
        qty = n * ct_mult
        pr = option_fill(action="open", bid=bid, ask=ask, qty=qty, fee_rate=fr)
        cost = pr.notional + pr.fee
        try:
            self.wallets.debit_trading(
                "USDC",
                cost,
                kind="option_open",
                note=f"buy {inst_id} x{n}@{pr.fill_px}",
            )
        except InsufficientFunds as e:
            return {"ok": False, "msg": str(e)}

        conn = self.get_db()
        try:
            existing = conn.execute(
                "SELECT * FROM sim_option_positions WHERE inst_id=?",
                (inst_id,),
            ).fetchone()
            if existing and float(existing["sheets"] or 0) > 1e-12:
                old_s = float(existing["sheets"])
                old_px = float(existing["entry_px"])
                old_prem = float(existing["premium_paid_usdc"] or 0)
                new_s = old_s + n
                entry = (old_px * old_s + pr.fill_px * n) / new_s
                conn.execute(
                    """
                    UPDATE sim_option_positions
                    SET sheets=?, entry_px=?, ct_mult=?, premium_paid_usdc=?, updated_at=?
                    WHERE inst_id=?
                    """,
                    (new_s, entry, ct_mult, old_prem + pr.notional, _now(), inst_id),
                )
            else:
                conn.execute(
                    """
                    INSERT INTO sim_option_positions(
                        inst_id, side, sheets, entry_px, ct_mult, premium_paid_usdc, updated_at
                    ) VALUES (?,?,?,?,?,?,?)
                    ON CONFLICT(inst_id) DO UPDATE SET
                        sheets=excluded.sheets,
                        entry_px=excluded.entry_px,
                        ct_mult=excluded.ct_mult,
                        premium_paid_usdc=excluded.premium_paid_usdc,
                        updated_at=excluded.updated_at
                    """,
                    (inst_id, "buy", float(n), pr.fill_px, ct_mult, pr.notional, _now()),
                )
            conn.commit()
        finally:
            conn.close()

        ord_id = f"sim-opt-{uuid.uuid4().hex[:16]}"
        self._store_option_order(
            ord_id=ord_id, inst_id=inst_id, side="buy", sheets=n, avg_px=pr.fill_px
        )
        return {
            "ok": True,
            "data": {"ordId": ord_id, "sCode": "0", "sMsg": "sim filled"},
            "raw": {"sim": True},
            "px": pr.fill_px,
            "ord_type": "ioc",
            "info": {"sim": True, "fee": pr.fee, "fill": pr.to_dict()},
        }

    def sell_option_close(
        self,
        exchange_options: Any,
        *,
        inst_id: str,
        sheets: int,
        price: float | None = None,
        fee_rate: float | None = None,
        **_kwargs: Any,
    ) -> dict[str, Any]:
        n = int(sheets)
        if n < 1:
            return {"ok": False, "msg": "张数至少为 1"}
        conn = self.get_db()
        try:
            row = conn.execute(
                "SELECT * FROM sim_option_positions WHERE inst_id=?",
                (inst_id,),
            ).fetchone()
            if not row or float(row["sheets"] or 0) <= 1e-12:
                return {"ok": False, "msg": "模拟期权无对应持仓"}
            pos_s = float(row["sheets"])
            ct_mult = float(row["ct_mult"] or 0.01)
            close_n = min(pos_s, float(n))
            if close_n <= 0:
                return {"ok": False, "msg": "平仓张数无效"}
            fr = sim_fee_rate(fee_rate)
            bid, ask, _ = _option_bid_ask(exchange_options, inst_id)
            if price is not None and float(price) > 0:
                bid = max(float(bid), float(price)) if float(price) > 0 else float(bid)
            qty = close_n * ct_mult
            pr = option_fill(action="close", bid=bid, ask=ask, qty=qty, fee_rate=fr)
            credit = pr.notional - pr.fee
            remain = pos_s - close_n
            if remain <= 1e-12:
                conn.execute("DELETE FROM sim_option_positions WHERE inst_id=?", (inst_id,))
            else:
                prem = float(row["premium_paid_usdc"] or 0) * (remain / pos_s)
                conn.execute(
                    """
                    UPDATE sim_option_positions
                    SET sheets=?, premium_paid_usdc=?, updated_at=?
                    WHERE inst_id=?
                    """,
                    (remain, prem, _now(), inst_id),
                )
            conn.commit()
        finally:
            conn.close()

        if credit > 0:
            self.wallets.credit_trading(
                "USDC",
                credit,
                kind="option_close",
                note=f"sell {inst_id} x{close_n}@{pr.fill_px}",
            )
        ord_id = f"sim-opt-close-{uuid.uuid4().hex[:16]}"
        self._store_option_order(
            ord_id=ord_id, inst_id=inst_id, side="sell", sheets=close_n, avg_px=pr.fill_px
        )
        return {
            "ok": True,
            "data": {"ordId": ord_id, "sCode": "0", "sMsg": "sim filled"},
            "raw": {"sim": True},
            "px": pr.fill_px,
            "ord_type": "ioc",
            "info": {"sim": True, "fee": pr.fee, "fill": pr.to_dict()},
        }

    def option_positions_okx_rows(self) -> list[dict[str, Any]]:
        """对齐 OKX positions 行字段, 供 format_position_row 使用."""
        rows = []
        for p in self.list_option_positions():
            rows.append(
                {
                    "instId": p["inst_id"],
                    "pos": str(p["sheets"]),
                    "avgPx": str(p["entry_px"]),
                    "markPx": str(p["entry_px"]),
                    "upl": "0",
                    "uplRatio": "0",
                    "posSide": "long",
                    "mgnMode": "isolated",
                }
            )
        return rows