crypto_monitor/lib/account_ledger/account_ledger_register.py

"""三所统一:账户流水路由 + 后台同步安装."""
from __future__ import annotations

import os
from typing import Any, Callable

from flask import Flask, Response, jsonify, request, session, stream_with_context
from jinja2 import ChoiceLoader, FileSystemLoader

from lib.account_ledger.account_ledger_db import ensure_account_ledger_tables, query_entries
from lib.account_ledger.account_ledger_normalize import (
    ACCOUNT_FUNDING,
    ACCOUNT_TRADING,
    PAGE_SIZE,
    VALID_ACCOUNTS,
)
from lib.account_ledger.account_ledger_sync import account_ledger_store
from lib.common.history_window_lib import resolve_list_window


def attach_account_ledger_templates(app: Flask, repo_root: str) -> None:
    tpl_dir = os.path.join(repo_root, "lib", "account_ledger", "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 _build_fetch_fn(exchange_key: str, app_module: Any) -> Callable:
    ex_key = (exchange_key or "").strip().lower()
    exchange = getattr(app_module, "exchange", None)
    ensure_markets = getattr(app_module, "ensure_markets_loaded", None)

    def _fetch(*, start_ms: int, end_ms: int):
        if exchange is None:
            return [], ["exchange missing"]
        if ex_key == "okx":
            from lib.exchange.okx_ledger_lib import fetch_okx_account_ledger

            return fetch_okx_account_ledger(
                exchange,
                start_ms=start_ms,
                end_ms=end_ms,
                ensure_markets=ensure_markets,
            )
        if ex_key == "binance":
            from lib.exchange.binance_ledger_lib import fetch_binance_account_ledger

            return fetch_binance_account_ledger(
                exchange,
                start_ms=start_ms,
                end_ms=end_ms,
                ensure_markets=ensure_markets,
            )
        from lib.exchange.gate_ledger_lib import fetch_gate_account_ledger

        return fetch_gate_account_ledger(
            exchange,
            start_ms=start_ms,
            end_ms=end_ms,
            ensure_markets=ensure_markets,
        )

    return _fetch


def _currencies_for_exchange(exchange_key: str) -> list[str]:
    if (exchange_key or "").strip().lower() == "okx":
        return ["USDT", "USDC"]
    return ["USDT"]


def install_account_ledger(
    app: Flask,
    repo_root: str,
    app_module: Any,
    *,
    exchange_key: str = "",
) -> None:
    ex = (exchange_key or "").strip().lower()
    if not ex:
        mod_name = getattr(app_module, "__name__", "") or ""
        if "okx" in mod_name.lower():
            ex = "okx"
        elif "binance" in mod_name.lower():
            ex = "binance"
        else:
            ex = "gate"
    exchange_key = ex

    attach_account_ledger_templates(app, repo_root)
    get_db = app_module.get_db
    login_required = app_module.login_required

    # 初始化表
    try:
        conn = get_db()
        try:
            ensure_account_ledger_tables(conn)
        finally:
            conn.close()
    except Exception:
        pass

    account_ledger_store.configure(
        get_db=get_db,
        fetch_fn=_build_fetch_fn(exchange_key, app_module),
        exchange_key=str(exchange_key),
    )
    account_ledger_store.start()
    app.extensions["account_ledger_exchange"] = str(exchange_key).lower()

    def _list_window():
        resolve = getattr(app_module, "_list_window_from_request", None)
        if callable(resolve):
            return resolve()
        return resolve_list_window(request.args, session)

    @app.route("/api/account_ledger")
    @login_required
    def api_account_ledger():
        account = (request.args.get("account") or ACCOUNT_FUNDING).strip().lower()
        if account not in VALID_ACCOUNTS:
            account = ACCOUNT_FUNDING
        try:
            page = int(request.args.get("page") or 1)
        except Exception:
            page = 1
        win = _list_window()
        start_ms = int(win.get("start_ms") or 0)
        end_ms = int(win.get("end_ms") or 0)
        ccys = _currencies_for_exchange(app.extensions.get("account_ledger_exchange") or "")
        conn = get_db()
        try:
            ensure_account_ledger_tables(conn)
            data = query_entries(
                conn,
                account=account,
                start_ms=start_ms,
                end_ms=end_ms,
                page=page,
                page_size=PAGE_SIZE,
                currencies=ccys,
            )
        finally:
            conn.close()
        st = account_ledger_store.status_dict()
        return jsonify(
            {
                "ok": True,
                "account": account,
                "window": {
                    "preset": win.get("preset"),
                    "label": win.get("label"),
                    "start_ms": start_ms,
                    "end_ms": end_ms,
                },
                "currencies": ccys,
                **data,
                **st,
            }
        )

    @app.route("/api/account_ledger/stream")
    @login_required
    def api_account_ledger_stream():
        return Response(
            stream_with_context(account_ledger_store.iter_sse()),
            mimetype="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "X-Accel-Buffering": "no",
            },
        )

    @app.route("/api/account_ledger/refresh", methods=["POST"])
    @login_required
    def api_account_ledger_refresh():
        win = _list_window()
        body = request.get_json(silent=True) or {}
        start_ms = body.get("start_ms", win.get("start_ms"))
        end_ms = body.get("end_ms", win.get("end_ms"))
        try:
            start_i = int(start_ms) if start_ms is not None else None
            end_i = int(end_ms) if end_ms is not None else None
        except Exception:
            start_i, end_i = None, None
        result = account_ledger_store.sync_once(
            reason="manual", start_ms=start_i, end_ms=end_i
        )
        return jsonify(result)

    @app.route("/account_ledger")
    @login_required
    def account_ledger_page():
        from lib.instance.instance_embed_lib import redirect_to_embed_shell_if_enabled

        redir = redirect_to_embed_shell_if_enabled("account_ledger")
        if redir is not None:
            return redir
        return app_module.render_main_page("account_ledger")