siloqy/prod/clean_arch/dita_v2/asex_account.py

"""ASEx wrapper for AccountProjectionV2 — serializes the 4 P0 accumulator sites."""
from __future__ import annotations

import time
from concurrent.futures import Future
from typing import Any, Iterable, List, Optional

from asex.guarded import ASExGuardedState
from asex.worker import ASExWorker

from .account import AccountProjectionV2, AccountSnapshotV2, EPosition, ReconcileConfig, TradeSlot


class _AccountBackend(ASExGuardedState[dict, Any]):
    """Wraps AccountProjectionV2 mutations as ASEx operations."""

    def __init__(self, seed_capital: float, **kw):
        super().__init__()
        self._proj = AccountProjectionV2(seed_capital, **kw)

    def _validate(self, mutation: dict) -> bool:
        return isinstance(mutation, dict) and "op" in mutation

    def _apply(self, mutation: dict) -> Any:
        op = mutation["op"]
        args = mutation.get("args", {})
        if op == "apply_fill":
            self._proj.apply_fill(**args); return None
        elif op == "apply_funding":
            self._proj.apply_funding(**args); return None
        elif op == "apply_balance_update":
            self._proj.apply_balance_update(**args); return None
        elif op == "apply_position_update":
            self._proj.apply_position_update(**args); return None
        elif op == "build_snapshot":
            return self._proj.build_snapshot(**args)
        raise ValueError(f"Unknown ASEx account op: {op}")

    def __getattr__(self, name):
        return getattr(self._proj, name)


class ASEXAccountV2:
    """ASEx-serialized wrapper around ``AccountProjectionV2``.

    All mutations go through an ``ASExWorker`` — one thread, sequential,
    no races.  Property reads bypass the worker (lock-free, fast).
    """

    def __init__(self, seed_capital: float, *, min_capital: float = 0.0,
                 max_capital: float | None = None,
                 reconcile_config: ReconcileConfig | None = None):
        self._backend = _AccountBackend(seed_capital, min_capital=min_capital,
                                         max_capital=max_capital, reconcile_config=reconcile_config)
        self._worker: ASExWorker = ASExWorker(self._backend, daemon=True)

    def apply_fill_async(self, **kw) -> Future:
        return self._worker.mutate({"op": "apply_fill", "args": kw})

    def apply_funding_async(self, amount: float) -> Future:
        return self._worker.mutate({"op": "apply_funding", "args": {"amount": amount}})

    def apply_balance_update_async(self, **kw) -> Future:
        return self._worker.mutate({"op": "apply_balance_update", "args": kw})

    def apply_position_update_async(self, positions: List[EPosition]) -> Future:
        return self._worker.mutate({"op": "apply_position_update", "args": {"positions": positions}})

    def build_snapshot_async(self, **kw) -> Future:
        return self._worker.mutate({"op": "build_snapshot", "args": kw})

    def apply_fill(self, **kw) -> None:
        self.apply_fill_async(**kw).result(timeout=30)

    def apply_funding(self, amount: float) -> None:
        self.apply_funding_async(amount).result(timeout=30)

    def apply_balance_update(self, **kw) -> None:
        self.apply_balance_update_async(**kw).result(timeout=30)

    def apply_position_update(self, positions: List[EPosition]) -> None:
        self.apply_position_update_async(positions).result(timeout=30)

    def build_snapshot(self, **kw) -> AccountSnapshotV2:
        return self.build_snapshot_async(**kw).result(timeout=30)

    @property
    def snapshot(self) -> AccountSnapshotV2:
        return self._backend._proj.snapshot

    @property
    def k_capital(self) -> float:
        return self._backend._proj.k_capital

    @property
    def applied(self) -> int:
        return self._backend.applied

    def close(self, *, timeout: float | None = None) -> None:
        self._worker.close(timeout=timeout)

    def __enter__(self) -> ASEXAccountV2:
        return self

    def __exit__(self, *args) -> None:
        self.close(timeout=5.0)