src/services/billing/gate.py

"""The request gate — authenticate, balance-check, and burn.

Every /v1/* request flows through here:
  1. resolve the Bearer token → (user_id | admin | None)
  2. if a user: check balance > 0, else 402
  3. after the upstream responds: compute cost from usage → spend from cache
     + record usage + ledger entry (async, best-effort, never blocks the response)

The master key bypasses billing entirely (admin tier).
"""
from __future__ import annotations

import hashlib
import logging
import re
from typing import Optional

from . import cache, pricing, registry, user_store

logger = logging.getLogger("router.gate")

# a valid 20-byte ethereum address: 0x + 40 hex chars (case-insensitive)
_ADDR_RE = re.compile(r"^0x[0-9a-f]{40}$")


def is_enabled() -> bool:
    """True if the billing system is fully configured (gist store + wallet).

    Used throughout to gracefully degrade to unmetered mode when billing
    secrets aren't set (e.g. local dev, or a Space without GIST_TOKEN).
    """
    from .. import gist_client, wallet  # local import avoids a cycle at module load
    return gist_client.is_configured() and wallet.is_configured()


def _sha256_hex(s: str) -> str:
    """sha256 hex digest of a string (how api keys are stored/compared)."""
    return hashlib.sha256(s.encode()).hexdigest()


def _normalize_addr(token: str) -> Optional[str]:
    """If `token` is an ETH-address key, return the lowercased address, else None.
    Accepts either `sk-0x…` (our key format) or a bare `0x…` address.
    """
    t = token.strip()
    if t[:3].lower() == "sk-":  # strip the sk- key prefix if present
        t = t[3:]
    t = t.strip().lower()
    return t if _ADDR_RE.match(t) else None  # valid address or None


def resolve(authorization: Optional[str], master_key: str) -> tuple[Optional[str], bool]:
    """Resolve a Bearer header to (user_id | None, is_admin).

    Returns (None, True) for the master key, (user_id, False) for a valid user
    credential, (None, False) for an invalid/missing token.

    A user credential is their ETH deposit address (as `sk-0x…` or bare `0x…`).
    Legacy `dr_u_…` keys still resolve via the sha256 index (backward compat).
    """
    if not authorization or not authorization.startswith("Bearer "):
        return None, False
    token = authorization[7:].strip()
    if not token:
        return None, False
    if master_key and token == master_key:  # admin tier — unmetered
        return None, True
    addr = _normalize_addr(token)  # address-key path (the primary credential)
    if addr:
        uid = registry.find_user_by_address(addr)  # O(1) address -> user lookup
        if uid:
            return uid, False
    key_hash = _sha256_hex(token)  # legacy dr_u_ path (sha256 -> user)
    user_id = registry.find_user_by_key_hash(key_hash)
    return user_id, False


async def check_balance(user_id: Optional[str]) -> bool:
    """True if the caller may proceed: admin (None) or a user with credits."""
    if user_id is None:  # admin
        return True
    return await cache.get_balance(user_id) > 0


async def burn(user_id: str, models: dict, model_name: str, tok_in: int, tok_out: int) -> int:
    """Charge a user for a request: spend credits + log usage + ledger. Returns credits charged.

    `free:` models cost 0 credits (pricing returns 0), but we STILL record the
    usage event + ledger entry so the dashboard reflects free-tier activity.
    """
    cost = pricing.cost_credits(models, model_name, tok_in, tok_out)
    _touch_active(user_id)  # stamp last-active (so the sweeper keeps this account)
    if cost > 0:  # paid model — spend credits now
        cache.spend(user_id, cost)  # decrement the in-memory balance instantly
    gid = cache.gist_id_for(user_id)  # fire-and-forget persistence (best-effort)
    if gid:
        try:
            await user_store.record_usage(gid, model_name, tok_in, tok_out, cost)  # counters + recent ring
            await user_store.append_ledger(gid, {
                "id": f"b_{tok_out}_{cost}",  # best-effort unique id
                "type": "burn",
                "model": model_name,
                "tok_in": tok_in,
                "tok_out": tok_out,
                "cost": cost,  # 0 for free models, >0 for paid
                "ts": registry._now(),
            })
        except Exception as e:  # logging failures must never break the request
            logger.warning("usage log failed for %s: %s", user_id, e)
    return cost


def _touch_active(user_id: str) -> None:
    """Stamp the current epoch on the user's in-memory registry record (for the sweeper)."""
    import time  # stdlib epoch
    rec = registry.find_user(user_id)  # the cached record
    if rec:  # found
        rec["last_active_epoch"] = time.time()  # update in-memory (persisted by the sweeper)


def extract_usage(payload: dict) -> tuple[int, int]:
    """Pull (prompt_tokens, completion_tokens) from an OpenAI-shaped response."""
    u = payload.get("usage") or {}
    return int(u.get("prompt_tokens", 0)), int(u.get("completion_tokens", 0))