crypto_okx/lib/sim/wallets_lib.py

"""模拟资金钱包: funding/trading × USDT/USDC."""

from __future__ import annotations

from datetime import datetime
from typing import Any, Callable


WALLET_KEYS = (
    "funding_usdt",
    "trading_usdt",
    "funding_usdc",
    "trading_usdc",
)

_ACCT_MAP = {
    ("funding", "usdt"): "funding_usdt",
    ("trading", "usdt"): "trading_usdt",
    ("funding", "usdc"): "funding_usdc",
    ("trading", "usdc"): "trading_usdc",
}


def normalize_sim_account(account: str | None) -> str:
    """统一账户别名: swap/unified → trading; spot → funding."""
    a = (account or "").strip().lower()
    if a in ("trading", "swap", "unified", "18"):
        return "trading"
    if a in ("funding", "spot", "6"):
        return "funding"
    return a


class InsufficientFunds(RuntimeError):
    pass


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

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

    def snapshot(self) -> dict[str, float]:
        conn = self.get_db()
        try:
            row = conn.execute("SELECT * FROM sim_wallets WHERE id=1").fetchone()
            if row is None:
                return {k: 0.0 for k in WALLET_KEYS}
            return {k: float(row[k] or 0) for k in WALLET_KEYS}
        finally:
            conn.close()

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

    def total_usdt_equiv(self, snap: dict[str, float] | None = None) -> float:
        v = snap or self.view()
        return round(
            float(v.get("funding_usdt") or 0)
            + float(v.get("trading_usdt") or 0)
            + float(v.get("funding_usdc") or 0)
            + float(v.get("trading_usdc") or 0),
            8,
        )

    def _write(self, snap: dict[str, float], conn=None) -> dict[str, float]:
        owns = conn is None
        if owns:
            conn = self.get_db()
        try:
            now = self._now()
            conn.execute(
                """
                UPDATE sim_wallets SET
                    funding_usdt=?, trading_usdt=?, funding_usdc=?, trading_usdc=?, updated_at=?
                WHERE id=1
                """,
                (
                    float(snap["funding_usdt"]),
                    float(snap["trading_usdt"]),
                    float(snap["funding_usdc"]),
                    float(snap["trading_usdc"]),
                    now,
                ),
            )
            if owns:
                conn.commit()
            return {k: float(snap[k]) for k in WALLET_KEYS}
        finally:
            if owns:
                conn.close()

    def _ledger(
        self,
        conn,
        *,
        kind: str,
        amount: float,
        ccy: str,
        account: str,
        balance_after: float,
        note: str = "",
    ) -> None:
        conn.execute(
            """
            INSERT INTO sim_ledger_entries(kind, amount, ccy, account, balance_after, note, ts)
            VALUES (?,?,?,?,?,?,?)
            """,
            (kind, float(amount), ccy.upper(), account, float(balance_after), note or "", self._now()),
        )

    def debit_trading(self, ccy: str, amount: float, *, kind: str = "debit", note: str = "") -> dict[str, float]:
        amt = float(amount)
        if amt <= 0:
            raise ValueError("扣款金额须大于 0")
        key = _ACCT_MAP.get(("trading", (ccy or "").lower()))
        if not key:
            raise ValueError("币种须为 USDT 或 USDC")
        conn = self.get_db()
        try:
            row = conn.execute("SELECT * FROM sim_wallets WHERE id=1").fetchone()
            snap = {k: float(row[k] or 0) for k in WALLET_KEYS}
            bal = float(snap[key])
            if amt > bal + 1e-9:
                raise InsufficientFunds(f"交易账户 {ccy.upper()} 不足(可用 {bal:.4f})")
            snap[key] = bal - amt
            self._write(snap, conn=conn)
            self._ledger(
                conn,
                kind=kind,
                amount=-amt,
                ccy=ccy,
                account="trading",
                balance_after=snap[key],
                note=note,
            )
            conn.commit()
            return snap
        finally:
            conn.close()

    def credit_trading(self, ccy: str, amount: float, *, kind: str = "credit", note: str = "") -> dict[str, float]:
        amt = float(amount)
        if amt < 0:
            raise ValueError("入账金额不能为负")
        if amt < 1e-12:
            return self.snapshot()
        key = _ACCT_MAP.get(("trading", (ccy or "").lower()))
        if not key:
            raise ValueError("币种须为 USDT 或 USDC")
        conn = self.get_db()
        try:
            row = conn.execute("SELECT * FROM sim_wallets WHERE id=1").fetchone()
            snap = {k: float(row[k] or 0) for k in WALLET_KEYS}
            snap[key] = float(snap[key]) + amt
            self._write(snap, conn=conn)
            self._ledger(
                conn,
                kind=kind,
                amount=amt,
                ccy=ccy,
                account="trading",
                balance_after=snap[key],
                note=note,
            )
            conn.commit()
            return snap
        finally:
            conn.close()

    def transfer(
        self,
        *,
        ccy: str,
        amount: float,
        from_account: str,
        to_account: str,
    ) -> dict[str, Any]:
        amt = float(amount)
        if amt <= 0:
            return {"ok": False, "detail": "划转金额须大于 0"}
        ccy_l = (ccy or "USDT").strip().lower()
        fa = normalize_sim_account(from_account)
        ta = normalize_sim_account(to_account)
        if fa not in ("funding", "trading") or ta not in ("funding", "trading"):
            return {"ok": False, "detail": "账户仅支持 funding / trading"}
        if fa == ta:
            return {"ok": False, "detail": "来源与目标账户不能相同"}
        src_key = _ACCT_MAP.get((fa, ccy_l))
        dst_key = _ACCT_MAP.get((ta, ccy_l))
        if not src_key or not dst_key:
            return {"ok": False, "detail": "币种须为 USDT 或 USDC"}
        conn = self.get_db()
        try:
            row = conn.execute("SELECT * FROM sim_wallets WHERE id=1").fetchone()
            snap = {k: float(row[k] or 0) for k in WALLET_KEYS}
            src_bal = float(snap[src_key])
            if amt > src_bal + 1e-9:
                return {"ok": False, "detail": f"余额不足(可用 {src_bal:.4f})"}
            snap[src_key] = src_bal - amt
            snap[dst_key] = float(snap[dst_key]) + amt
            self._write(snap, conn=conn)
            self._ledger(
                conn,
                kind="transfer",
                amount=-amt,
                ccy=ccy_l,
                account=fa,
                balance_after=snap[src_key],
                note=f"to {ta}",
            )
            self._ledger(
                conn,
                kind="transfer",
                amount=amt,
                ccy=ccy_l,
                account=ta,
                balance_after=snap[dst_key],
                note=f"from {fa}",
            )
            conn.commit()
            return {
                "ok": True,
                "detail": "transferred",
                "ccy": ccy_l.upper(),
                "amount": amt,
                "from": fa,
                "to": ta,
                "wallets": {k: float(snap[k]) for k in WALLET_KEYS},
                "total_usdt_equiv": self.total_usdt_equiv(snap),
            }
        finally:
            conn.close()

    def convert(
        self,
        *,
        from_ccy: str,
        to_ccy: str,
        amount: float,
        account: str = "trading",
        to_amount: float | None = None,
        rate: float | None = None,
        fee: float | None = None,
        note: str | None = None,
    ) -> dict[str, Any]:
        """USDT↔USDC 兑换. 默认交易账户; to_amount 未给时按 rate(USDT/USDC) 换算, 再否则 1:1."""
        amt = float(amount)
        if amt <= 0:
            return {"ok": False, "detail": "数量须大于 0"}
        fa = (from_ccy or "").strip().lower()
        ta = (to_ccy or "").strip().lower()
        acct = normalize_sim_account(account) or "trading"
        if acct not in ("funding", "trading"):
            return {"ok": False, "detail": "account 须为 funding / trading"}
        if {fa, ta} != {"usdt", "usdc"}:
            return {"ok": False, "detail": "仅支持 USDT↔USDC"}
        if to_amount is not None:
            got = float(to_amount)
        elif rate is not None and float(rate) > 0:
            r = float(rate)
            # rate = USDT per 1 USDC
            got = (amt / r) if fa == "usdt" else (amt * r)
        else:
            got = amt
        if got <= 0:
            return {"ok": False, "detail": "兑换所得须大于 0"}
        src_key = _ACCT_MAP[(acct, fa)]
        dst_key = _ACCT_MAP[(acct, ta)]
        conn = self.get_db()
        try:
            row = conn.execute("SELECT * FROM sim_wallets WHERE id=1").fetchone()
            snap = {k: float(row[k] or 0) for k in WALLET_KEYS}
            src = float(snap[src_key])
            if amt > src + 1e-9:
                return {"ok": False, "detail": f"{acct} {fa.upper()} 不足(可用 {src:.4f})"}
            snap[src_key] = src - amt
            snap[dst_key] = float(snap[dst_key]) + got
            self._write(snap, conn=conn)
            note_s = note or f"to {ta} @{rate if rate is not None else '1:1'}"
            self._ledger(
                conn,
                kind="convert",
                amount=-amt,
                ccy=fa,
                account=acct,
                balance_after=snap[src_key],
                note=note_s,
            )
            self._ledger(
                conn,
                kind="convert",
                amount=got,
                ccy=ta,
                account=acct,
                balance_after=snap[dst_key],
                note=f"from {fa}",
            )
            conn.commit()
            eff_rate = (amt / got) if fa == "usdt" and got > 0 else ((got / amt) if amt > 0 else None)
            return {
                "ok": True,
                "detail": "converted",
                "from_ccy": fa.upper(),
                "to_ccy": ta.upper(),
                "amount": amt,
                "to_amount": got,
                "rate": float(rate) if rate is not None else eff_rate,
                "fee": float(fee or 0),
                "account": acct,
                "wallets": {k: float(snap[k]) for k in WALLET_KEYS},
                "total_usdt_equiv": self.total_usdt_equiv(snap),
            }
        finally:
            conn.close()

    def has_open_positions(self) -> bool:
        conn = self.get_db()
        try:
            p = conn.execute(
                "SELECT COUNT(*) AS n FROM sim_perp_positions WHERE contracts > 1e-12"
            ).fetchone()
            o = conn.execute(
                "SELECT COUNT(*) AS n FROM sim_option_positions WHERE sheets > 1e-12"
            ).fetchone()
            pn = int(p["n"] if hasattr(p, "keys") else p[0])
            on = int(o["n"] if hasattr(o, "keys") else o[0])
            return pn > 0 or on > 0
        finally:
            conn.close()

    def reset_equity(self, amount: float, *, force: bool = False) -> dict[str, Any]:
        if self.has_open_positions() and not force:
            return {"ok": False, "detail": "仍有模拟持仓,请先平仓或传 force=true"}
        amt = max(0.0, float(amount))
        conn = self.get_db()
        try:
            if force:
                conn.execute("DELETE FROM sim_perp_positions")
                conn.execute("DELETE FROM sim_option_positions")
                conn.execute("DELETE FROM sim_option_orders")
            now = self._now()
            snap = {
                "funding_usdt": amt,
                "trading_usdt": 0.0,
                "funding_usdc": 0.0,
                "trading_usdc": 0.0,
            }
            self._write(snap, conn=conn)
            self._ledger(
                conn,
                kind="reset",
                amount=amt,
                ccy="USDT",
                account="funding",
                balance_after=amt,
                note="reset equity",
            )
            conn.commit()
            return {"ok": True, "wallets": snap, "total_usdt_equiv": amt}
        finally:
            conn.close()