app/services/token_budget.py

"""Author RAG Chatbot SaaS — Token budget helpers.

Single source of truth for budget math and post-chat usage updates.
"""

from app.models.client_access import ClientAccess

BUDGET_EXHAUSTED_PREFIX = "budget_exhausted:"
TOKENS_REDIS_PREFIX = "tokens:"


def total_budget(access: ClientAccess) -> int:
    """Total available tokens including bonus grants."""
    return (access.token_budget or 0) + (access.bonus_tokens or 0)


def tokens_remaining(access: ClientAccess) -> int:
    """Tokens left in the current billing period."""
    return max(0, total_budget(access) - (access.tokens_used or 0))


def percent_used(access: ClientAccess) -> float:
    """Percentage of total budget consumed."""
    budget = total_budget(access)
    if budget <= 0:
        return 0.0
    return round((access.tokens_used or 0) / budget * 100, 1)


def usage_summary(access: ClientAccess) -> dict:
    """Standard token usage dict for API responses."""
    budget = total_budget(access)
    used = access.tokens_used or 0
    return {
        "budget": budget,
        "used": used,
        "remaining": max(0, budget - used),
        "percent_used": percent_used(access),
        "has_active_subscription": True,
        "subscription_status": "active",
        "plan": getattr(access, "plan", None),
        "expires_at": getattr(access, "expires_at", None) and access.expires_at.isoformat(),
    }


NO_SUBSCRIPTION_SUMMARY: dict = {
    "budget": 0,
    "used": 0,
    "remaining": 0,
    "percent_used": 0,
    "has_active_subscription": False,
    "subscription_status": "none",
    "plan": None,
    "expires_at": None,
}


def usage_summary_or_empty(access: ClientAccess | None) -> dict:
    """Return usage stats for an active grant, or zeros when unsubscribed."""
    if not access:
        return dict(NO_SUBSCRIPTION_SUMMARY)
    return usage_summary(access)


async def record_token_usage(
    db,
    redis,
    author_id: str,
    token_count: int,
) -> None:
    """Increment usage counters and set budget_exhausted when at 100%.

    Updates client_access.tokens_used (DB) and tokens:{author_id}:current (Redis).
    """
    if token_count <= 0:
        return

    from app.repositories.access_repo import AccessRepository

    token_key = f"{TOKENS_REDIS_PREFIX}{author_id}:current"
    await redis.incrby(token_key, token_count)
    await redis.expire(token_key, 32 * 24 * 3600)

    access_repo = AccessRepository(db)
    access = await access_repo.get_active_for_author(author_id)
    if not access:
        return

    # R-014/R-052: Atomic SQL UPDATE instead of Python read-modify-write
    from sqlalchemy import update
    from app.models.client_access import ClientAccess as CA
    await db.execute(
        update(CA)
        .where(CA.id == access.id, CA.author_id == access.author_id)
        .values(tokens_used=CA.tokens_used + token_count)
    )
    await db.flush()
    await db.refresh(access)
    budget = total_budget(access)

    if budget > 0 and access.tokens_used >= budget:
        await redis.set(f"{BUDGET_EXHAUSTED_PREFIX}{author_id}", "1", ex=32 * 24 * 3600)


async def clear_budget_exhausted(redis, author_id: str) -> None:
    """Clear the budget exhausted flag (e.g. after bonus tokens)."""
    await redis.delete(f"{BUDGET_EXHAUSTED_PREFIX}{author_id}")


async def clear_token_usage_counters(redis, author_id: str) -> None:
    """Clear Redis usage counters after SuperAdmin reset."""
    await redis.delete(
        f"{BUDGET_EXHAUSTED_PREFIX}{author_id}",
        f"{TOKENS_REDIS_PREFIX}{author_id}:current",
    )


# ── Budget Warning (B8 fix) ───────────────────────────────────────────────────

_BUDGET_WARN_PREFIX = "budget_warned:"


async def check_and_warn_budget(redis, access: ClientAccess, author) -> None:
    """Send a one-time warning email when an author hits 80% token usage.

    B8 fix: Previously authors had zero warning before their chatbot stopped
    working. This function checks usage after each chat turn and sends a single
    email when the threshold is crossed. Redis prevents duplicate emails within
    the same grant period.

    Gracefully skips if:
    - Usage is below threshold
    - Email was already sent this grant period
    - Email service is not configured (SMTP not set)

    Args:
        redis: Async Redis connection.
        access: Active ClientAccess grant object.
        author: Author User object (for email address).
    """
    from app.config import get_settings
    cfg = get_settings()

    pct = percent_used(access)
    if pct < cfg.TOKEN_WARNING_THRESHOLD_PCT * 100:
        return

    warn_key = f"{_BUDGET_WARN_PREFIX}{access.id}"
    already_warned = await redis.get(warn_key)
    if already_warned:
        return  # Already sent warning this grant period

    if not getattr(author, "notify_token_alerts", True):
        return

    try:
        from datetime import datetime, timezone
        from app.services.email_service import EmailService

        budget = total_budget(access)
        used = access.tokens_used or 0
        remaining = tokens_remaining(access)
        days_left = 30
        if access.expires_at:
            days_left = max((access.expires_at - datetime.now(timezone.utc)).days, 1)
        daily = used / max(30 - days_left, 1) if used else 0
        forecast = "soon" if daily > 0 and remaining / daily < 7 else f"in ~{days_left} days"

        email_svc = EmailService()
        email_svc.send_token_budget_warning(
            to=author.email,
            name=author.full_name or author.email,
            pct_used=int(round(pct, 0)),
            tokens_used=used,
            tokens_total=budget,
            forecast_date=forecast,
        )
    except Exception as e:
        import structlog
        structlog.get_logger(__name__).warning(
            "Budget warning email failed — skipping (non-fatal)", error=str(e)
        )
        return

    from datetime import datetime, timezone
    remaining_secs = 86400
    if access.expires_at:
        delta = (access.expires_at - datetime.now(timezone.utc)).total_seconds()
        remaining_secs = max(3600, int(delta))
    await redis.setex(warn_key, remaining_secs, "1")

    try:
        from app.tasks.webhook_delivery_task import deliver_webhook_event

        deliver_webhook_event.delay(
            author.id,
            "low_budget_warning",
            {
                "percent_used": round(pct, 1),
                "tokens_used": access.tokens_used or 0,
                "tokens_total": budget,
            },
        )
    except Exception:
        pass