exec-mcp/src/exec_mcp/services_bitstamp.py

from __future__ import annotations

import os
import threading

try:
    import bitstamp.client
except ModuleNotFoundError:  # allows test runs without the optional dependency
    bitstamp = None  # type: ignore

from . import repo
from .bitstamp import BitstampClient
from .bitstamp_fx import load_eur_usd

BITSTAMP_BASE_URL = "https://www.bitstamp.net"
# Live trading reads should refresh quickly, but not so often that Bitstamp trips a breaker.
BALANCE_CACHE_TTL_MS = 1200
ACCOUNT_INFO_CACHE_TTL_MS = 1200
FEES_CACHE_TTL_SECONDS = 10 * 60
STALE_CACHE_TTL_SECONDS = 2
_CACHE_LOCKS: dict[str, threading.Lock] = {}
_CACHE_LOCKS_GUARD = threading.Lock()
_BITSTAMP_CLIENTS: dict[str, BitstampClient] = {}
_BITSTAMP_CLIENTS_GUARD = threading.Lock()


def _ttl_from_env_ms(name: str, default_ms: int) -> float:
    try:
        raw = os.getenv(name)
        if raw is not None:
            return max(float(raw) / 1000.0, 0.0)
        return max(float(default_ms) / 1000.0, 0.0)
    except Exception:
        return max(float(default_ms) / 1000.0, 0.0)


def _cache_lock(cache_key: str) -> threading.Lock:
    with _CACHE_LOCKS_GUARD:
        lock = _CACHE_LOCKS.get(cache_key)
        if lock is None:
            lock = threading.Lock()
            _CACHE_LOCKS[cache_key] = lock
        return lock


def _cache_error(cache_key: str, detail: str, ttl_seconds: int = 15) -> None:
    repo.cache_put(cache_key, {"_cached_error": detail}, ttl_seconds)


def _raise_cached_error(payload: dict) -> None:
    detail = str(payload.get("_cached_error") or "Bitstamp request failed")
    raise RuntimeError(detail)


def _stale_key(cache_key: str) -> str:
    return f"{cache_key}:stale"


def invalidate_account_cache(account_id: str) -> None:
    for key in (
        f"bitstamp:account_balance:{account_id}",
        f"bitstamp:account_balance:{account_id}:stale",
        f"bitstamp:account_info:{account_id}",
        f"bitstamp:account_info:{account_id}:stale",
    ):
        repo.cache_delete(key)


def invalidate_balance_cache(account_id: str) -> None:
    for key in (
        f"bitstamp:account_balance:{account_id}",
        f"bitstamp:account_balance:{account_id}:stale",
    ):
        repo.cache_delete(key)


def invalidate_account_info_cache(account_id: str) -> None:
    for key in (
        f"bitstamp:account_info:{account_id}",
        f"bitstamp:account_info:{account_id}:stale",
    ):
        repo.cache_delete(key)


def _fee_cache_key(market_symbol: str | None = None) -> str:
    return "bitstamp:fees:all" if market_symbol is None else f"bitstamp:fees:{market_symbol.lower()}"


def _normalize_fee_entry(item: dict) -> dict:
    fees = item.get("fees") or {}
    return {
        "currency_pair": item.get("currency_pair"),
        "market": item.get("market"),
        "fees": {
            "maker": fees.get("maker"),
            "taker": fees.get("taker"),
        },
    }


def fetch_trading_fees(account_id: str, market_symbol: str | None = None) -> dict:
    cache_key = _fee_cache_key(market_symbol)
    cached = repo.cache_get(cache_key)
    if cached is not None:
        return cached

    with _cache_lock(cache_key):
        cached = repo.cache_get(cache_key)
        if cached is not None:
            return cached

        try:
            client = get_bitstamp_client(account_id)
            if market_symbol:
                payload = client.trading.fees_trading_market(market_symbol)
            else:
                payload = client.trading.fees_trading()
        except Exception as exc:
            _cache_error(cache_key, str(exc))
            raise

        if isinstance(payload, list):
            result = {"source": "bitstamp", "market": market_symbol, "fees": [_normalize_fee_entry(item) for item in payload]}
        elif isinstance(payload, dict):
            result = {"source": "bitstamp", "market": market_symbol, **_normalize_fee_entry(payload)}
        else:
            result = {"source": "bitstamp", "market": market_symbol, "raw": payload}

        repo.cache_put(cache_key, result, FEES_CACHE_TTL_SECONDS)
        return result


def _require_client() -> None:
    if bitstamp is None:
        raise RuntimeError("bitstamp-python-client dependency is not installed")


def _build_trading_client(account_id: str):
    _require_client()
    account = repo.get_account(account_id)
    secrets = repo.get_account_secrets(account_id)
    return bitstamp.client.Trading(
        username=account["venue_account_ref"],
        key=secrets["api_key"],
        secret=secrets["api_secret"],
    )


def get_bitstamp_client(account_id: str):
    with _BITSTAMP_CLIENTS_GUARD:
        client = _BITSTAMP_CLIENTS.get(account_id)
        if client is None:
            account = repo.get_account(account_id)
            secrets = repo.get_account_secrets(account_id)
            client = BitstampClient(
                username=account["venue_account_ref"],
                api_key=secrets["api_key"],
                api_secret=secrets["api_secret"],
            )
            _BITSTAMP_CLIENTS[account_id] = client
        return client


def clear_bitstamp_trading_client(account_id: str) -> None:
    with _BITSTAMP_CLIENTS_GUARD:
        _BITSTAMP_CLIENTS.pop(account_id, None)


def _normalize_account_balances_payload(payload: list[dict], account_id: str) -> list[dict]:
    balances: list[dict] = []
    for item in payload:
        currency = item.get("currency")
        if not currency:
            continue
        try:
            total = float(item.get("total", 0) or 0)
            available = float(item.get("available", 0) or 0)
            reserved = float(item.get("reserved", 0) or 0)
        except (TypeError, ValueError):
            continue
        balances.append(
            {
                "account_id": account_id,
                "asset_code": str(currency).upper(),
                "available": available,
                "reserved": reserved,
                "total": total,
            }
        )
    return balances


def _price_for_asset(asset_code: str) -> float | None:
    asset = str(asset_code or "").lower()
    if not asset:
        return None
    if asset == "usd":
        return 1.0
    if asset == "eur":
        fx = load_eur_usd()
        if fx and fx.get("sell") is not None:
            try:
                return float(fx["sell"])
            except (TypeError, ValueError):
                return None
        return None
    return repo.get_latest_price(f"{asset}usd")


def _value_account_balances(balance: dict) -> tuple[list[dict], float]:
    valued_balances: list[dict] = []
    total_value_usd = 0.0
    for item in balance.get("balances", []):
        asset_code = str(item.get("asset_code") or "").upper()
        total = float(item.get("total") or 0)
        price_usd = _price_for_asset(asset_code)
        value_usd = total * price_usd if price_usd is not None else None
        if value_usd is not None:
            total_value_usd += value_usd
        valued_balances.append(
            {
                **item,
                "asset_code": asset_code,
                "price_usd": price_usd,
                "value_currency": "USD",
                "value_usd": value_usd,
            }
        )
    return valued_balances, total_value_usd


def fetch_account_balance(account_id: str) -> dict:
    cache_key = f"bitstamp:account_balance:{account_id}"
    cached = repo.cache_get(cache_key)
    if cached is not None:
        if isinstance(cached, dict) and cached.get("_cached_error"):
            stale = repo.cache_get(_stale_key(cache_key))
            if stale is not None:
                return stale
            _raise_cached_error(cached)
        return cached

    with _cache_lock(cache_key):
        cached = repo.cache_get(cache_key)
        if cached is not None:
            if isinstance(cached, dict) and cached.get("_cached_error"):
                stale = repo.cache_get(_stale_key(cache_key))
                if stale is not None:
                    return stale
                _raise_cached_error(cached)
            return cached

        try:
            client = get_bitstamp_client(account_id)
            payload = client.trading._post("account_balances/", return_json=True, version=2)
        except Exception as exc:
            _cache_error(cache_key, str(exc))
            raise

        normalized = _normalize_account_balances_payload(payload, account_id)

        result = {"source": "bitstamp", "cached": False, "balances": normalized, "payload": payload}
        repo.cache_put(
            cache_key,
            result,
            _ttl_from_env_ms("BITSTAMP_BALANCE_CACHE_TTL_MS", BALANCE_CACHE_TTL_MS),
        )
        repo.cache_put(_stale_key(cache_key), result, STALE_CACHE_TTL_SECONDS)
        return result


def capture_account_balance_snapshot(account_id: str, *, source_kind: str, source_ref: str = "") -> dict:
    if source_ref and repo.has_account_balance_snapshot(account_id=account_id, source_kind=source_kind, source_ref=source_ref):
        return {"ok": True, "account_id": account_id, "source_kind": source_kind, "source_ref": source_ref, "snapshot_id": None, "snapshot": None}

    account = repo.get_account(account_id)
    invalidate_balance_cache(account_id)
    balance = fetch_account_balance(account_id)
    valued_balances, total_value_usd = _value_account_balances(balance)
    snapshot = {
        "account_id": account["id"],
        "display_name": account["display_name"],
        "venue": account["venue"],
        "venue_account_ref": account["venue_account_ref"],
        "balances": valued_balances,
        "total_value_usd": total_value_usd,
        "raw_balance": balance["payload"],
    }
    snapshot_id = repo.save_account_balance_snapshot(
        account_id=account_id,
        source_kind=source_kind,
        source_ref=source_ref,
        snapshot=snapshot,
    )
    return {
        "ok": True,
        "account_id": account_id,
        "source_kind": source_kind,
        "source_ref": source_ref,
        "snapshot_id": snapshot_id,
        "snapshot": snapshot,
    }


def fetch_account_info(account_id: str) -> dict:
    cache_key = f"bitstamp:account_info:{account_id}"
    cached = repo.cache_get(cache_key)
    if cached is not None:
        if isinstance(cached, dict) and cached.get("_cached_error"):
            stale = repo.cache_get(_stale_key(cache_key))
            if stale is not None:
                return stale
            _raise_cached_error(cached)
        return cached

    with _cache_lock(cache_key):
        cached = repo.cache_get(cache_key)
        if cached is not None:
            if isinstance(cached, dict) and cached.get("_cached_error"):
                stale = repo.cache_get(_stale_key(cache_key))
                if stale is not None:
                    return stale
                _raise_cached_error(cached)
            return cached

        try:
            account = repo.get_account(account_id)
            balance = fetch_account_balance(account_id)
        except Exception as exc:
            _cache_error(cache_key, str(exc))
            raise

        valued_balances, total_value_usd = _value_account_balances(balance)

        result = {
            "id": account["id"],
            "display_name": account["display_name"],
            "venue": account["venue"],
            "venue_account_ref": account["venue_account_ref"],
            "description": account["description"],
            "enabled": account["enabled"],
            "metadata": account["metadata"],
            "balances": valued_balances,
            "total_value_usd": total_value_usd,
            "raw_balance": balance["payload"],
        }

        repo.cache_put(
            cache_key,
            result,
            _ttl_from_env_ms("BITSTAMP_ACCOUNT_INFO_CACHE_TTL_MS", ACCOUNT_INFO_CACHE_TTL_MS),
        )
        repo.cache_put(_stale_key(cache_key), result, STALE_CACHE_TTL_SECONDS)
        return result


def fetch_account_fees(account_id: str, market_symbol: str | None = None) -> dict:
    return fetch_trading_fees(account_id, market_symbol)