app/control/account/state_machine.py

"""Account lifecycle state machine.

Applies feedback events to an AccountRecord, advancing its status
and updating quota / usage fields accordingly.
"""

from dataclasses import dataclass, field

from app.platform.runtime.clock import now_ms
from .enums import AccountStatus, FeedbackKind
from .models import AccountRecord, QuotaWindow
from .quota_defaults import default_quota_set

# ---------------------------------------------------------------------------
# Policy
# ---------------------------------------------------------------------------


@dataclass(slots=True)
class StatePolicy:
    """Configurable thresholds for automatic state transitions."""

    fail_threshold: int = 5  # consecutive failures → COOLING
    forbidden_strikes: int = 1  # 403 count → DISABLED
    default_cooling_ms: int = 15 * 60 * 1000  # 15 min


_DEFAULT_POLICY = StatePolicy()


# ---------------------------------------------------------------------------
# Feedback
# ---------------------------------------------------------------------------


@dataclass(slots=True)
class AccountFeedback:
    """Outcome of a single upstream request.

    ``mode_id``     — which quota window to update (0=auto, 1=fast, 2=expert).
    ``quota_window`` — real quota window from the usage API (optional).
    ``apply_usage`` — whether to increment usage_use_count on SUCCESS.
    """

    kind: FeedbackKind
    mode_id: int = 0
    at: int = field(default_factory=now_ms)
    status_code: int | None = None
    reason: str = ""
    quota_window: QuotaWindow | None = None  # real value from usage API
    retry_after_ms: int | None = None
    confirm_expired: bool = False
    apply_usage: bool = True

    @classmethod
    def from_status_code(
        cls,
        status_code: int,
        mode_id: int = 0,
        *,
        reason: str = "",
        retry_after_ms: int | None = None,
        confirm_expired: bool = False,
    ) -> "AccountFeedback":
        if status_code == 401:
            kind = FeedbackKind.UNAUTHORIZED
        elif status_code == 403:
            kind = FeedbackKind.FORBIDDEN
        elif status_code == 429:
            kind = FeedbackKind.RATE_LIMITED
        elif status_code >= 500:
            kind = FeedbackKind.SERVER_ERROR
        elif 200 <= status_code < 300:
            kind = FeedbackKind.SUCCESS
        else:
            kind = FeedbackKind.SERVER_ERROR
        return cls(
            kind=kind,
            mode_id=mode_id,
            status_code=status_code,
            reason=reason,
            retry_after_ms=retry_after_ms,
            confirm_expired=confirm_expired,
            apply_usage=False,
        )


# ---------------------------------------------------------------------------
# State helpers
# ---------------------------------------------------------------------------

_COOLDOWN_UNTIL_KEY = "cooldown_until"
_COOLDOWN_REASON_KEY = "cooldown_reason"
_DISABLED_AT_KEY = "disabled_at"
_DISABLED_REASON_KEY = "disabled_reason"
_EXPIRED_AT_KEY = "expired_at"
_EXPIRED_REASON_KEY = "expired_reason"
_FORBIDDEN_STRIKE_KEY = "forbidden_strikes"


def derive_status(record: AccountRecord, *, now: int | None = None) -> AccountStatus:
    """Compute the effective status, considering cooldown expiry."""
    if record.status != AccountStatus.COOLING:
        return record.status
    cooldown_until = record.ext.get(_COOLDOWN_UNTIL_KEY)
    if cooldown_until is None:
        return AccountStatus.COOLING
    ts = now if now is not None else now_ms()
    if ts >= int(cooldown_until):
        return AccountStatus.ACTIVE
    return AccountStatus.COOLING


def is_selectable(
    record: AccountRecord, mode_id: int, *, now: int | None = None
) -> bool:
    """Return True if the account can be selected for *mode_id*."""
    if record.is_deleted():
        return False
    status = derive_status(record, now=now)
    if status != AccountStatus.ACTIVE:
        return False
    qs = record.quota_set()
    win = qs.get(mode_id)
    return win is not None and not win.is_exhausted()


def is_manageable(record: AccountRecord, *, now: int | None = None) -> bool:
    """Return True if the account should participate in maintenance flows."""
    if record.is_deleted():
        return False
    status = derive_status(record, now=now)
    return status in (AccountStatus.ACTIVE, AccountStatus.COOLING)


# ---------------------------------------------------------------------------
# State transition
# ---------------------------------------------------------------------------


def apply_feedback(
    record: AccountRecord,
    feedback: AccountFeedback,
    *,
    policy: StatePolicy = _DEFAULT_POLICY,
) -> AccountRecord:
    """Apply *feedback* to *record* and return an updated copy.

    The returned object is a **new** record; the original is not mutated.
    """
    ts = feedback.at
    ext = dict(record.ext)
    qs = record.quota_set()

    status = record.status
    last_fail_at = record.last_fail_at
    last_fail_reason = record.last_fail_reason
    last_use_at = record.last_use_at
    last_sync_at = record.last_sync_at
    use_count = record.usage_use_count
    fail_count = record.usage_fail_count
    sync_count = record.usage_sync_count
    state_reason = record.state_reason

    # Update quota window from real API data if provided.
    if feedback.quota_window is not None:
        qs.set(feedback.mode_id, feedback.quota_window)
        last_sync_at = ts
        sync_count += 1
    elif feedback.kind == FeedbackKind.SUCCESS:
        # Decrement locally when no real data is available.
        win = qs.get(feedback.mode_id)
        if win is not None:
            updated_win = QuotaWindow(
                remaining=max(0, win.remaining - 1),
                total=win.total,
                window_seconds=win.window_seconds,
                reset_at=win.reset_at,
                synced_at=win.synced_at,
                source=win.source,
            )
            qs.set(feedback.mode_id, updated_win)
    elif feedback.kind == FeedbackKind.RATE_LIMITED:
        # Mark quota as zero; set reset_at from retry_after if available.
        win = qs.get(feedback.mode_id)
        if win is not None:
            reset_at = (
                ts + feedback.retry_after_ms
                if feedback.retry_after_ms
                else (ts + win.window_seconds * 1000)
            )
            qs.set(
                feedback.mode_id,
                QuotaWindow(
                    remaining=0,
                    total=win.total,
                    window_seconds=win.window_seconds,
                    reset_at=reset_at,
                    synced_at=win.synced_at,
                    source=win.source,
                ),
            )

    # Update usage counters.
    if feedback.kind == FeedbackKind.SUCCESS and feedback.apply_usage:
        use_count += 1
        last_use_at = ts
    elif feedback.kind not in (
        FeedbackKind.SUCCESS,
        FeedbackKind.RESTORE,
        FeedbackKind.DISABLE,
        FeedbackKind.DELETE,
    ):
        fail_count += 1
        last_fail_at = ts
        last_fail_reason = feedback.reason or str(feedback.status_code or "")

    # Status transitions.
    if feedback.kind == FeedbackKind.UNAUTHORIZED:
        if feedback.confirm_expired:
            status = AccountStatus.EXPIRED
            state_reason = feedback.reason or "token_expired"
            ext[_EXPIRED_AT_KEY] = ts
            ext[_EXPIRED_REASON_KEY] = state_reason
        else:
            # Unconfirmed 401 — only note failure; do not expire.
            pass

    elif feedback.kind == FeedbackKind.FORBIDDEN:
        strikes = int(ext.get(_FORBIDDEN_STRIKE_KEY, 0)) + 1
        ext[_FORBIDDEN_STRIKE_KEY] = strikes
        if strikes >= policy.forbidden_strikes:
            status = AccountStatus.DISABLED
            state_reason = feedback.reason or "forbidden"
            ext[_DISABLED_AT_KEY] = ts
            ext[_DISABLED_REASON_KEY] = state_reason

    elif feedback.kind == FeedbackKind.RATE_LIMITED:
        cooldown_ms = (
            feedback.retry_after_ms
            if feedback.retry_after_ms
            else policy.default_cooling_ms
        )
        status = AccountStatus.COOLING
        state_reason = feedback.reason or "rate_limited"
        ext[_COOLDOWN_UNTIL_KEY] = ts + cooldown_ms
        ext[_COOLDOWN_REASON_KEY] = state_reason

    elif feedback.kind == FeedbackKind.SUCCESS:
        if status == AccountStatus.COOLING:
            cooldown_until = ext.get(_COOLDOWN_UNTIL_KEY)
            if cooldown_until is not None and ts >= int(cooldown_until):
                status = AccountStatus.ACTIVE
                state_reason = None
                ext.pop(_COOLDOWN_UNTIL_KEY, None)
                ext.pop(_COOLDOWN_REASON_KEY, None)
                ext.pop(_FORBIDDEN_STRIKE_KEY, None)

    elif feedback.kind == FeedbackKind.DISABLE:
        status = AccountStatus.DISABLED
        state_reason = feedback.reason or "operator_disabled"
        ext[_DISABLED_AT_KEY] = ts
        ext[_DISABLED_REASON_KEY] = state_reason

    elif feedback.kind == FeedbackKind.DELETE:
        pass  # Caller sets deleted_at.

    elif feedback.kind == FeedbackKind.RESTORE:
        status = AccountStatus.ACTIVE
        state_reason = None
        ext.pop(_COOLDOWN_UNTIL_KEY, None)
        ext.pop(_COOLDOWN_REASON_KEY, None)
        ext.pop(_DISABLED_AT_KEY, None)
        ext.pop(_DISABLED_REASON_KEY, None)
        ext.pop(_EXPIRED_AT_KEY, None)
        ext.pop(_EXPIRED_REASON_KEY, None)
        ext.pop(_FORBIDDEN_STRIKE_KEY, None)
        # Reset quota to defaults.
        qs = default_quota_set(record.pool)

    return record.model_copy(
        update={
            "status": status,
            "quota": qs.to_dict(),
            "usage_use_count": use_count,
            "usage_fail_count": fail_count,
            "usage_sync_count": sync_count,
            "last_use_at": last_use_at,
            "last_fail_at": last_fail_at,
            "last_fail_reason": last_fail_reason,
            "last_sync_at": last_sync_at,
            "state_reason": state_reason,
            "ext": ext,
            "updated_at": ts,
        }
    )


def clear_failures(record: AccountRecord) -> AccountRecord:
    """Reset failure counters and restore ACTIVE status."""
    ext = dict(record.ext)
    for k in (
        _COOLDOWN_UNTIL_KEY,
        _COOLDOWN_REASON_KEY,
        _DISABLED_AT_KEY,
        _DISABLED_REASON_KEY,
        _EXPIRED_AT_KEY,
        _EXPIRED_REASON_KEY,
        _FORBIDDEN_STRIKE_KEY,
    ):
        ext.pop(k, None)
    return record.model_copy(
        update={
            "status": AccountStatus.ACTIVE,
            "usage_fail_count": 0,
            "last_fail_at": None,
            "last_fail_reason": None,
            "state_reason": None,
            "ext": ext,
            "updated_at": now_ms(),
        }
    )


__all__ = [
    "StatePolicy",
    "AccountFeedback",
    "derive_status",
    "is_selectable",
    "is_manageable",
    "apply_feedback",
    "clear_failures",
]