app/dataplane/account/__init__.py

"""AccountDirectory — high-concurrency hot-path account store.

Wraps the columnar AccountRuntimeTable with lock-minimal coordination.
Bootstrap loads a full snapshot; incremental sync applies revision-based
changesets without holding the selection lock.
"""

import asyncio
from typing import TYPE_CHECKING

from app.platform.config.snapshot import get_config
from app.platform.logging.logger import logger
from app.platform.runtime.clock import now_s
from app.control.account.repository import AccountRepository
from app.control.account.enums import FeedbackKind
from .table import AccountRuntimeTable
from .lease import AccountLease, new_lease
from .selector import current_strategy, select, select_any
from .sync import bootstrap as _bootstrap, apply_changes
from . import feedback as fb
from ..shared.enums import POOL_ID_TO_STR, StatusId

if TYPE_CHECKING:
    pass


class AccountDirectory:
    """High-concurrency, lock-minimal account store for the hot path.

    Design:
      - ``_table`` is replaced atomically; readers never see partial state.
      - ``_lock`` guards all table mutations (feedback, sync, status updates).
      - Selection (``reserve``) copies the candidate set under lock, then
        scores outside the lock — lock is re-acquired only to increment
        inflight.
    """

    def __init__(self, repository: AccountRepository) -> None:
        self._repo = repository
        self._table: AccountRuntimeTable | None = None
        self._lock = asyncio.Lock()
        self._sync_lock = asyncio.Lock()

    # ------------------------------------------------------------------
    # Lifecycle
    # ------------------------------------------------------------------

    async def bootstrap(self) -> None:
        """Load initial snapshot from the repository."""
        table = await _bootstrap(self._repo)
        async with self._lock:
            self._table = table
        logger.info("account directory ready: size={}", table.size)

    async def sync_if_changed(self) -> bool:
        """Apply incremental changes since last revision.

        Returns True if any changes were applied.
        Uses a dedicated sync lock so only one sync runs at a time,
        without blocking the selection lock.
        """
        if self._table is None:
            return False

        async with self._sync_lock:
            async with self._lock:
                table = self._table
            changed = await apply_changes(table, self._repo)
            if changed:
                logger.debug(
                    "account directory synced: revision={} size={}",
                    table.revision,
                    table.size,
                )
            return changed

    # ------------------------------------------------------------------
    # Selection (hot path)
    # ------------------------------------------------------------------

    async def reserve(
        self,
        pool_candidates: tuple[int, ...] | int,
        mode_id: int,
        *,
        exclude_tokens: list[str] | None = None,
        prefer_tags: list[str] | None = None,
        now_s_override: int | None = None,
    ) -> AccountLease | None:
        """Select and reserve the best available account slot.

        ``pool_candidates`` is tried in order; the first pool with an available
        account wins. A plain ``int`` is accepted and wrapped into a tuple.

        Returns an AccountLease, or None if no account is available.
        """
        table = self._table
        if table is None:
            return None

        pools: tuple[int, ...] = (
            (pool_candidates,) if isinstance(pool_candidates, int) else pool_candidates
        )
        ts = now_s_override if now_s_override is not None else now_s()

        # Resolve exclude set (O(n) once, before lock).
        exclude_idxs: frozenset[int] | None = None
        if exclude_tokens:
            idxs = [
                table.idx_by_token[t] for t in exclude_tokens if t in table.idx_by_token
            ]
            if idxs:
                exclude_idxs = frozenset(idxs)

        # Resolve tag preference set.
        prefer_tag_idxs: set[int] | None = None
        if prefer_tags:
            sets = [
                table.tag_idx.get(tag) for tag in prefer_tags if tag in table.tag_idx
            ]
            if sets:
                prefer_tag_idxs = set().union(*sets)

        async with self._lock:
            idx: int | None = None
            for pool_id in pools:
                idx = select(
                    table,
                    pool_id,
                    mode_id,
                    exclude_idxs=exclude_idxs,
                    prefer_tag_idxs=prefer_tag_idxs,
                    now_s=ts,
                )
                if idx is not None:
                    break

            if idx is None:
                return None

            fb.increment_inflight(table, idx)
            fb.update_last_use(table, idx, ts)
            token = table.get_token(idx)
            actual_pool = table.get_pool_id(idx)

        return new_lease(
            idx=idx,
            token=token,
            pool_id=actual_pool,
            mode_id=mode_id,
            selected_at=ts,
        )

    async def reserve_any(
        self,
        pool_candidates: tuple[int, ...] | int,
        *,
        exclude_tokens: list[str] | None = None,
        prefer_tags: list[str] | None = None,
        now_s_override: int | None = None,
    ) -> AccountLease | None:
        """Select any active account from pool_candidates without mode quota checking.

        Used for WebSocket-based operations that manage their own upstream rate limiting.
        Returns an AccountLease with mode_id=-1 (no specific mode is tracked).
        """
        table = self._table
        if table is None:
            return None

        pools: tuple[int, ...] = (
            (pool_candidates,) if isinstance(pool_candidates, int) else pool_candidates
        )
        ts = now_s_override if now_s_override is not None else now_s()

        exclude_idxs: frozenset[int] | None = None
        if exclude_tokens:
            idxs = [
                table.idx_by_token[t] for t in exclude_tokens if t in table.idx_by_token
            ]
            if idxs:
                exclude_idxs = frozenset(idxs)

        prefer_tag_idxs: set[int] | None = None
        if prefer_tags:
            sets = [
                table.tag_idx.get(tag) for tag in prefer_tags if tag in table.tag_idx
            ]
            if sets:
                prefer_tag_idxs = set().union(*sets)

        async with self._lock:
            idx: int | None = None
            for pool_id in pools:
                idx = select_any(
                    table,
                    pool_id,
                    exclude_idxs=exclude_idxs,
                    prefer_tag_idxs=prefer_tag_idxs,
                    now_s=ts,
                )
                if idx is not None:
                    break

            if idx is None:
                return None

            fb.increment_inflight(table, idx)
            fb.update_last_use(table, idx, ts)
            token = table.get_token(idx)
            actual_pool = table.get_pool_id(idx)

        return new_lease(
            idx=idx,
            token=token,
            pool_id=actual_pool,
            mode_id=-1,  # no specific mode tracked for WS operations
            selected_at=ts,
        )

    async def release(self, lease: AccountLease) -> None:
        """Decrement inflight counter for a finished request."""
        table = self._table
        if table is None:
            return
        async with self._lock:
            fb.decrement_inflight(table, lease.idx)

    # ------------------------------------------------------------------
    # Feedback (hot path)
    # ------------------------------------------------------------------

    async def feedback(
        self,
        token: str,
        kind: FeedbackKind,
        mode_id: int,
        *,
        remaining: int | None = None,
        reset_at_ms: int | None = None,
        now_s_val: int | None = None,
    ) -> None:
        """Apply upstream response feedback to the account slot."""
        table = self._table
        if table is None:
            return

        idx = table.idx_by_token.get(token)
        if idx is None:
            return

        ts = now_s_val if now_s_val is not None else now_s()

        strategy = current_strategy()

        async with self._lock:
            if kind == FeedbackKind.SUCCESS:
                if strategy == "random":
                    fb.apply_success_random(table, idx)
                else:
                    fb.apply_success_quota(table, idx, mode_id)

            elif kind == FeedbackKind.RATE_LIMITED:
                if strategy == "random":
                    pool_id = int(table.pool_by_idx[idx])
                    cooling_sec = _pool_cooling_sec(pool_id)
                    fb.apply_rate_limited_random(table, idx, cooling_sec=cooling_sec)
                else:
                    fb.apply_rate_limited_quota(table, idx, mode_id)
                fb.update_last_fail(table, idx, ts)

            elif kind == FeedbackKind.UNAUTHORIZED:
                fb.apply_auth_failure(table, idx)
                fb.update_last_fail(table, idx, ts)
                fb.apply_status_change(table, idx, int(StatusId.EXPIRED))

            elif kind == FeedbackKind.FORBIDDEN:
                fb.apply_forbidden(table, idx)
                fb.update_last_fail(table, idx, ts)

            elif kind == FeedbackKind.SERVER_ERROR:
                fb.apply_server_error(table, idx)
                fb.update_last_fail(table, idx, ts)

            # Quota strategy may receive authoritative quota data from upstream
            # response headers; the random strategy ignores this entirely.
            if strategy == "quota" and remaining is not None and reset_at_ms is not None:
                reset_s = int(reset_at_ms // 1000)
                fb.apply_quota_update(table, idx, mode_id, remaining, reset_s)

    # ------------------------------------------------------------------
    # Diagnostics
    # ------------------------------------------------------------------

    @property
    def size(self) -> int:
        return self._table.size if self._table else 0

    @property
    def revision(self) -> int:
        return self._table.revision if self._table else 0


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


_POOL_INTERVAL_CONFIG: dict[str, tuple[str, int]] = {
    "basic": ("account.refresh.basic_interval_sec", 86_400),
    "super": ("account.refresh.super_interval_sec", 7_200),
    "heavy": ("account.refresh.heavy_interval_sec", 7_200),
}


def _pool_cooling_sec(pool_id: int) -> int:
    """Cooling seconds for a 429 on a given pool (random strategy only)."""
    pool_str = POOL_ID_TO_STR.get(pool_id, "basic")
    interval_key, default_interval = _POOL_INTERVAL_CONFIG.get(
        pool_str, _POOL_INTERVAL_CONFIG["basic"]
    )
    return max(0, int(get_config(interval_key, default_interval)))


# ---------------------------------------------------------------------------
# Module-level singleton
# ---------------------------------------------------------------------------

_directory: AccountDirectory | None = None


async def get_account_directory(
    repository: AccountRepository | None = None,
) -> AccountDirectory:
    """Return the module-level AccountDirectory, bootstrapping on first call.

    After first initialization, ``repository`` may be omitted.
    """
    global _directory
    if _directory is None:
        if repository is None:
            raise RuntimeError(
                "AccountDirectory not bootstrapped — repository required on first call"
            )
        _directory = AccountDirectory(repository)
        await _directory.bootstrap()
    return _directory


__all__ = ["AccountDirectory", "get_account_directory"]