chat_backend/llm_be/finance/services/quotas.py

"""Prompt-window and token-period quota checks (shared by chat + finance APIs)."""

from __future__ import annotations

from dataclasses import dataclass
from datetime import timedelta
from typing import Any

from django.conf import settings
from django.db.models import Count, Q, Sum
from django.utils import timezone

from chat_backend.models import PromptMetric
from finance.models import UserSubscription
from finance.services.plans import get_or_create_user_subscription, seed_subscription_plans


class QuotaExceeded(Exception):
    """Raised when a generation turn is blocked by quota."""

    def __init__(self, code: str, message: str, *, details: dict | None = None):
        super().__init__(message)
        self.code = code
        self.message = message
        self.details = details or {}


class FeatureNotAllowed(Exception):
    """Raised when the user's plan cannot use a feature."""

    def __init__(self, code: str, message: str, *, details: dict | None = None):
        super().__init__(message)
        self.code = code
        self.message = message
        self.details = details or {}


@dataclass
class UsageSnapshot:
    prompts_in_window: int
    prompt_quota: int | None
    prompts_remaining: int | None
    window_hours: int
    tokens_in_period: int | None
    tokens_out_period: int | None
    tokens_total_period: int | None
    turns_missing_token_usage: int
    monthly_token_quota: int | None
    tokens_remaining: int | None
    period_start: Any
    period_end: Any

    def to_dict(self) -> dict[str, Any]:
        return {
            "prompts_in_window": self.prompts_in_window,
            "prompt_quota": self.prompt_quota,
            "prompts_remaining": self.prompts_remaining,
            "window_hours": self.window_hours,
            "tokens_in_period": self.tokens_in_period,
            "tokens_out_period": self.tokens_out_period,
            "tokens_total_period": self.tokens_total_period,
            "turns_missing_token_usage": self.turns_missing_token_usage,
            "monthly_token_quota": self.monthly_token_quota,
            "tokens_remaining": self.tokens_remaining,
            "period_start": self.period_start.isoformat() if self.period_start else None,
            "period_end": self.period_end.isoformat() if self.period_end else None,
        }


def _user_conversation_ids(user) -> list[int]:
    from chat_backend.models import Conversation

    return list(
        Conversation.objects.filter(user=user, deleted=False).values_list("id", flat=True)
    )


def _billing_period_bounds():
    """Calendar-month UTC window for token-period aggregation (#17)."""
    now = timezone.now()
    start = now.replace(day=1, hour=0, minute=0, second=0, microsecond=0)
    if start.month == 12:
        end = start.replace(year=start.year + 1, month=1)
    else:
        end = start.replace(month=start.month + 1)
    return start, end


def _sum_tokens(qs) -> tuple[int | None, int | None]:
    """
    Sum tokens_in / tokens_out.

    Returns (None, None) when *no* rows reported usage — never fabricate 0.
    When some rows reported usage, sum only those (nulls ignored by Sum).
    """
    agg = qs.aggregate(
        tin=Sum("tokens_in"),
        tout=Sum("tokens_out"),
        with_in=Count("id", filter=Q(tokens_in__isnull=False)),
        with_out=Count("id", filter=Q(tokens_out__isnull=False)),
    )
    tokens_in = agg["tin"] if agg["with_in"] else None
    tokens_out = agg["tout"] if agg["with_out"] else None
    return tokens_in, tokens_out


def get_usage_snapshot(user) -> UsageSnapshot:
    seed_subscription_plans(update_existing=False)
    sub = (
        UserSubscription.objects.select_related("plan")
        .filter(user_id=user.pk)
        .first()
    )
    if sub is None:
        sub = get_or_create_user_subscription(user)
    plan = sub.plan if sub.is_active else None

    window_hours = plan.prompt_window_hours if plan else 6
    prompt_quota = plan.prompt_quota_per_window if plan else None
    monthly_token_quota = sub.effective_monthly_token_quota() if sub.is_active else None

    conversation_ids = _user_conversation_ids(user)
    now = timezone.now()
    window_start = now - timedelta(hours=window_hours)
    period_start, period_end = _billing_period_bounds()

    base = PromptMetric.objects.filter(conversation_id__in=conversation_ids)

    prompts_in_window = base.filter(created__gte=window_start).count()
    period_qs = base.filter(created__gte=period_start, created__lt=period_end)
    tokens_in, tokens_out = _sum_tokens(period_qs)
    missing = period_qs.filter(
        Q(tokens_in__isnull=True) | Q(tokens_out__isnull=True)
    ).count()

    if tokens_in is None and tokens_out is None:
        tokens_total = None
    else:
        tokens_total = (tokens_in or 0) + (tokens_out or 0)

    prompts_remaining = None
    if prompt_quota is not None:
        prompts_remaining = max(prompt_quota - prompts_in_window, 0)

    tokens_remaining = None
    if monthly_token_quota is not None and tokens_total is not None:
        tokens_remaining = max(monthly_token_quota - tokens_total, 0)
    elif monthly_token_quota is not None and tokens_total is None:
        # No provider usage yet — do not treat as 0 consumed.
        tokens_remaining = monthly_token_quota

    return UsageSnapshot(
        prompts_in_window=prompts_in_window,
        prompt_quota=prompt_quota,
        prompts_remaining=prompts_remaining,
        window_hours=window_hours,
        tokens_in_period=tokens_in,
        tokens_out_period=tokens_out,
        tokens_total_period=tokens_total,
        turns_missing_token_usage=missing,
        monthly_token_quota=monthly_token_quota,
        tokens_remaining=tokens_remaining,
        period_start=period_start,
        period_end=period_end,
    )


def assert_feature_allowed(user, feature: str) -> None:
    if not getattr(settings, "ENFORCE_SUBSCRIPTION_GATES", True):
        return

    seed_subscription_plans(update_existing=False)
    sub = (
        UserSubscription.objects.select_related("plan")
        .filter(user_id=user.pk)
        .first()
    )

    if sub is None or not sub.is_active or sub.plan is None:
        raise FeatureNotAllowed(
            "subscription_required",
            "An active subscription is required to use this feature.",
            details={"feature": feature},
        )

    if not sub.plan.allows_feature(feature):
        raise FeatureNotAllowed(
            "feature_not_allowed",
            f"Your plan ({sub.plan.name}) does not include {feature.replace('_', ' ')}.",
            details={
                "feature": feature,
                "plan": sub.plan.slug,
            },
        )


def assert_within_quotas(user) -> UsageSnapshot:
    """
    Enforce prompt-window (#36) and token-period (#17) limits.

    Precedence: either limit may block. Missing provider token usage does not
    silently under-count toward a token cap — turns with null tokens are tracked
    in `turns_missing_token_usage` and token-cap enforcement only uses reported
    sums; if quota is set and usage is entirely unknown, we allow the turn but
    surface the gap (callers/admin can tighten later).
    """
    seed_subscription_plans(update_existing=False)
    sub = (
        UserSubscription.objects.select_related("plan")
        .filter(user_id=user.pk)
        .first()
    )

    if sub is None or not sub.is_active or sub.plan is None:
        raise QuotaExceeded(
            "subscription_required",
            "An active subscription is required before sending prompts.",
        )

    usage = get_usage_snapshot(user)

    if usage.prompt_quota is not None and usage.prompts_in_window >= usage.prompt_quota:
        raise QuotaExceeded(
            "prompt_quota_exceeded",
            (
                f"Prompt limit reached ({usage.prompt_quota} per "
                f"{usage.window_hours} hours). Try again later."
            ),
            details=usage.to_dict(),
        )

    if (
        usage.monthly_token_quota is not None
        and usage.tokens_total_period is not None
        and usage.tokens_total_period >= usage.monthly_token_quota
    ):
        raise QuotaExceeded(
            "token_quota_exceeded",
            (
                f"Monthly token limit reached ({usage.monthly_token_quota}). "
                "Upgrade or wait for the next billing period."
            ),
            details=usage.to_dict(),
        )

    return usage


def check_generation_allowed(user, *, feature: str = "text_generation") -> UsageSnapshot:
    """Combined feature + quota gate for a chat turn."""
    if not getattr(settings, "ENFORCE_SUBSCRIPTION_GATES", True):
        return get_usage_snapshot(user)
    assert_feature_allowed(user, feature)
    return assert_within_quotas(user)