crypto_okx/lib/sim/register.py

"""安装模拟资金路由与模板."""

from __future__ import annotations

import os
from typing import Any

from flask import Flask, jsonify, request
from jinja2 import ChoiceLoader, FileSystemLoader

from lib.sim.broker_lib import SimBroker
from lib.sim.db_lib import init_sim_tables
from lib.sim.hooks import apply_sim_hooks, patch_options_cfg, set_get_db
from lib.sim.mode_lib import get_trading_mode, is_sim_mode, set_trading_mode
from lib.sim.pricing_lib import sim_fee_rate
from lib.sim.wallets_lib import SimWallets


def attach_sim_templates(app: Flask, repo_root: str) -> None:
    tpl_dir = os.path.join(repo_root, "lib", "sim", "templates")
    if not os.path.isdir(tpl_dir):
        return
    existing = app.jinja_loader
    loaders = [FileSystemLoader(tpl_dir)]
    if existing is not None:
        if isinstance(existing, ChoiceLoader):
            loaders = list(existing.loaders) + loaders
        else:
            loaders.insert(0, existing)
    app.jinja_loader = ChoiceLoader(loaders)


def install_sim_trading(app: Flask, repo_root: str, app_module: Any = None) -> None:
    if app_module is None:
        raise ValueError("install_sim_trading 需要 app_module")

    attach_sim_templates(app, repo_root)
    get_db = app_module.get_db
    login_required = app_module.login_required
    set_get_db(get_db)

    try:
        conn = get_db()
        try:
            init_sim_tables(conn)
        finally:
            conn.close()
    except Exception as e:
        print(f"[sim] init tables: {e}")

    apply_sim_hooks(app_module)

    # 若 options / hedge 已安装, 补丁其 cfg, 并刷新永续下单引用
    for key in ("options_cfg", "hedge_plan_cfg"):
        cfg = app.extensions.get(key)
        if isinstance(cfg, dict):
            patch_options_cfg(cfg)
            for fn_name in (
                "place_exchange_order",
                "close_exchange_order",
                "get_exchange_capitals",
                "get_available_trading_usdt",
                "ensure_okx_live_ready",
                "get_live_position_contracts",
            ):
                if fn_name in cfg and hasattr(app_module, fn_name):
                    cfg[fn_name] = getattr(app_module, fn_name)

    app.extensions["sim_installed"] = True

    def _status_payload():
        mode = get_trading_mode(get_db)
        wallets = SimWallets(get_db).view() if mode == "sim" else None
        return {
            "mode": mode,
            "is_sim": mode == "sim",
            "wallets": wallets,
            "fee_rate": sim_fee_rate(),
        }

    @app.route("/api/sim/status")
    @login_required
    def api_sim_status():
        return jsonify({"ok": True, **_status_payload()})

    @app.route("/api/sim/mode", methods=["POST"])
    @login_required
    def api_sim_mode():
        body = request.get_json(silent=True) or {}
        mode = body.get("mode")
        try:
            saved = set_trading_mode(get_db, mode)
        except ValueError as e:
            return jsonify({"ok": False, "msg": str(e)}), 400
        return jsonify({"ok": True, "mode": saved, **_status_payload()})

    @app.route("/api/sim/reset", methods=["POST"])
    @login_required
    def api_sim_reset():
        if not is_sim_mode(get_db):
            return jsonify({"ok": False, "msg": "仅模拟模式可重置"}), 400
        body = request.get_json(silent=True) or {}
        try:
            equity = float(body.get("equity_usdt") if body.get("equity_usdt") is not None else 10000)
        except (TypeError, ValueError):
            return jsonify({"ok": False, "msg": "equity_usdt 无效"}), 400
        force = bool(body.get("force"))
        result = SimWallets(get_db).reset_equity(equity, force=force)
        if not result.get("ok"):
            return jsonify(result), 400
        return jsonify({"ok": True, **result, **_status_payload()})

    @app.route("/api/sim/transfer", methods=["POST"])
    @login_required
    def api_sim_transfer():
        if not is_sim_mode(get_db):
            return jsonify({"ok": False, "msg": "仅模拟模式可划转"}), 400
        body = request.get_json(silent=True) or {}
        try:
            amount = float(body.get("amount") or 0)
        except (TypeError, ValueError):
            return jsonify({"ok": False, "msg": "amount 无效"}), 400
        result = SimWallets(get_db).transfer(
            ccy=str(body.get("ccy") or "USDT"),
            amount=amount,
            from_account=str(body.get("from") or ""),
            to_account=str(body.get("to") or ""),
        )
        if not result.get("ok"):
            return jsonify(result), 400
        return jsonify({**result, **_status_payload()})

    @app.route("/api/sim/convert", methods=["POST"])
    @login_required
    def api_sim_convert():
        if not is_sim_mode(get_db):
            return jsonify({"ok": False, "msg": "仅模拟模式可兑换"}), 400
        body = request.get_json(silent=True) or {}
        try:
            amount = float(body.get("amount") or 0)
        except (TypeError, ValueError):
            return jsonify({"ok": False, "msg": "amount 无效"}), 400
        from_ccy = str(body.get("from_ccy") or "USDT").strip().upper()
        to_ccy = str(body.get("to_ccy") or "USDC").strip().upper()
        if {from_ccy, to_ccy} != {"USDT", "USDC"}:
            return jsonify({"ok": False, "msg": "仅支持 USDT↔USDC"}), 400
        direction = "usdt_to_usdc" if from_ccy == "USDT" else "usdc_to_usdt"
        account = str(body.get("account") or "trading")
        ex = getattr(app_module, "exchange", None) or getattr(app_module, "exchange_options", None)
        if ex is None:
            return jsonify({"ok": False, "msg": "无公开行情 exchange"}), 500
        try:
            result = SimBroker(get_db).convert_usdt_usdc(
                ex, direction=direction, amount=amount, account=account
            )
        except Exception as e:
            return jsonify({"ok": False, "msg": str(e)}), 400
        if not result.get("ok"):
            return jsonify(result), 400
        return jsonify({**result, **_status_payload()})

    @app.route("/api/sim/positions")
    @login_required
    def api_sim_positions():
        if not is_sim_mode(get_db):
            return jsonify({"ok": True, "perp": [], "options": [], "is_sim": False})
        b = SimBroker(get_db)
        return jsonify(
            {
                "ok": True,
                "is_sim": True,
                "perp": b.list_perp_positions(),
                "options": b.list_option_positions(),
            }
        )