chat_backend/llm_be/finance/services/plans.py

"""Subscription plan catalog helpers, seeding, and user assignment."""

from __future__ import annotations

import logging
from typing import Any

from django.db import transaction
from django.utils import timezone

from chat_backend.models import UserAuthEvent
from finance.models import BackerEmail, SubscriptionPlan, UserSubscription

logger = logging.getLogger(__name__)


def log_subscription_auth_event(
    user,
    *,
    started: bool,
    detail: str,
) -> None:
    event_type = (
        UserAuthEvent.EventType.SUBSCRIPTION_STARTED
        if started
        else UserAuthEvent.EventType.SUBSCRIPTION_UPDATED
    )
    UserAuthEvent.log(user, event_type, detail=detail[:512])

# Seed catalog for #36. Standard/Pro/Business stay hidden until explicitly enabled.
PLAN_SEED: list[dict[str, Any]] = [
    {
        "slug": SubscriptionPlan.Slug.FOUNDERS,
        "name": "Founders",
        "description": (
            "Unlimited product access for early supporters: text plus all future "
            "capabilities as they ship. $10/mo."
        ),
        "price_cents": 1000,
        "is_public": True,
        "is_selectable": True,
        "allows_text_generation": True,
        "allows_image_generation": True,
        "allows_rag": True,
        "allows_all_future_features": True,
        "prompt_quota_per_window": 300,
        "prompt_window_hours": 6,
        "monthly_token_quota": None,
        "sort_order": 10,
    },
    {
        "slug": SubscriptionPlan.Slug.STANDARD,
        "name": "Standard",
        "description": (
            "Secure conversational chat and coding assistance for developers "
            "and privacy-conscious individuals."
        ),
        "price_cents": 1500,
        "is_public": False,
        "is_selectable": False,
        "allows_text_generation": True,
        "allows_image_generation": False,
        "allows_rag": False,
        "allows_all_future_features": False,
        "prompt_quota_per_window": 100,
        "prompt_window_hours": 6,
        "monthly_token_quota": 1_000_000,
        "sort_order": 20,
    },
    {
        "slug": SubscriptionPlan.Slug.PRO,
        "name": "Pro / Creator",
        "description": (
            "Higher message caps and multi-modal workflows for heavy users, "
            "including image generation and Drive/RAG sync when available."
        ),
        "price_cents": 4000,
        "is_public": False,
        "is_selectable": False,
        "allows_text_generation": True,
        "allows_image_generation": True,
        "allows_rag": True,
        "allows_all_future_features": False,
        "prompt_quota_per_window": 200,
        "prompt_window_hours": 6,
        "monthly_token_quota": 3_000_000,
        "sort_order": 30,
    },
    {
        "slug": SubscriptionPlan.Slug.BUSINESS,
        "name": "Business Team",
        "description": (
            "Team seats, centralized auth, priority support, company Drive/RAG "
            "sync, and absolute data privacy for local companies handling "
            "sensitive data."
        ),
        "price_cents": 9900,
        "is_public": False,
        "is_selectable": False,
        "allows_text_generation": True,
        "allows_image_generation": True,
        "allows_rag": True,
        "allows_all_future_features": False,
        "prompt_quota_per_window": 300,
        "prompt_window_hours": 6,
        "monthly_token_quota": 5_000_000,
        "sort_order": 40,
    },
    {
        "slug": SubscriptionPlan.Slug.BACKER,
        "name": "Backer",
        "description": (
            "Complimentary Founders-level access for pre-approved emails. "
            "Not shown at checkout."
        ),
        "price_cents": 0,
        "is_public": False,
        "is_selectable": False,
        "allows_text_generation": True,
        "allows_image_generation": True,
        "allows_rag": True,
        "allows_all_future_features": True,
        "prompt_quota_per_window": 300,
        "prompt_window_hours": 6,
        "monthly_token_quota": None,
        "sort_order": 5,
    },
]


def seed_subscription_plans(*, update_existing: bool = True) -> list[SubscriptionPlan]:
    """Idempotently create/update the canonical plan catalog."""
    plans: list[SubscriptionPlan] = []
    for row in PLAN_SEED:
        slug = row["slug"]
        defaults = {k: v for k, v in row.items() if k != "slug"}
        plan, created = SubscriptionPlan.objects.get_or_create(
            slug=slug,
            defaults=defaults,
        )
        if not created and update_existing:
            for key, value in defaults.items():
                setattr(plan, key, value)
            plan.save()
        plans.append(plan)
    return plans


def get_plan(slug: str) -> SubscriptionPlan | None:
    return SubscriptionPlan.objects.filter(slug=slug).first()


def get_or_create_user_subscription(user) -> UserSubscription:
    sub, _ = UserSubscription.objects.get_or_create(user=user)
    return sub


def assign_plan(
    user,
    *,
    plan: SubscriptionPlan,
    source: str,
    status: str = UserSubscription.Status.ACTIVE,
    stripe_subscription_id: str = "",
    log_auth_event: bool = True,
) -> UserSubscription:
    sub = get_or_create_user_subscription(user)
    prev_plan_id = sub.plan_id
    prev_status = sub.status
    prev_source = sub.source
    prev_stripe_sub = sub.stripe_subscription_id or ""
    had_active = (
        prev_status == UserSubscription.Status.ACTIVE and prev_plan_id is not None
    )

    sub.plan = plan
    sub.source = source
    sub.status = status
    if stripe_subscription_id:
        sub.stripe_subscription_id = stripe_subscription_id
    sub.save()

    if log_auth_event:
        became_active = (
            status == UserSubscription.Status.ACTIVE and plan is not None
        )
        changed = (
            prev_plan_id != sub.plan_id
            or prev_status != sub.status
            or prev_source != sub.source
            or (
                bool(stripe_subscription_id)
                and prev_stripe_sub != (sub.stripe_subscription_id or "")
            )
        )
        if became_active and not had_active:
            log_subscription_auth_event(
                user,
                started=True,
                detail=(
                    f"plan={plan.slug} source={source} status={status}"
                    + (
                        f" stripe_subscription_id={stripe_subscription_id}"
                        if stripe_subscription_id
                        else ""
                    )
                ),
            )
        elif changed:
            log_subscription_auth_event(
                user,
                started=False,
                detail=(
                    f"plan={plan.slug} source={source} status={status}"
                    + (
                        f" stripe_subscription_id={stripe_subscription_id}"
                        if stripe_subscription_id
                        else ""
                    )
                ),
            )
    return sub


def user_has_active_plan(user) -> bool:
    try:
        sub = user.subscription
    except UserSubscription.DoesNotExist:
        return False
    return sub.is_active


def needs_checkout(user) -> bool:
    """True when the user must complete paid Checkout to use the product."""
    try:
        sub = user.subscription
    except UserSubscription.DoesNotExist:
        return True
    if not sub.is_active:
        return True
    # Complimentary / already-paid tiers skip Checkout.
    if sub.source in (
        UserSubscription.Source.BACKER,
        UserSubscription.Source.ADMIN,
        UserSubscription.Source.STRIPE,
    ):
        return False
    return True


@transaction.atomic
def try_redeem_backer_email(user) -> UserSubscription | None:
    """
    If the user's email is on the Backer whitelist and unused, assign Backer plan.

    Returns the UserSubscription when redeemed, else None.
    """
    email = (getattr(user, "email", "") or "").strip().lower()
    if not email:
        return None

    entry = (
        BackerEmail.objects.select_for_update()
        .filter(email__iexact=email, redeemed_at__isnull=True)
        .first()
    )
    if entry is None:
        return None

    seed_subscription_plans(update_existing=False)
    plan = get_plan(SubscriptionPlan.Slug.BACKER)
    if plan is None:
        logger.error("Backer plan missing from catalog; cannot redeem %s", email)
        return None

    sub = assign_plan(
        user,
        plan=plan,
        source=UserSubscription.Source.BACKER,
        status=UserSubscription.Status.ACTIVE,
    )
    entry.redeemed_at = timezone.now()
    entry.redeemed_user = user
    entry.save(update_fields=["redeemed_at", "redeemed_user", "last_modified"])
    logger.info("Redeemed Backer email %s for user %s", email, user.pk)
    return sub


def assign_founders_from_stripe(
    user,
    *,
    stripe_subscription_id: str = "",
) -> UserSubscription:
    """Backward-compatible helper; prefer ``assign_plan_from_stripe``."""
    return assign_plan_from_stripe(
        user,
        plan_slug=SubscriptionPlan.Slug.FOUNDERS,
        stripe_subscription_id=stripe_subscription_id,
    )


def assign_plan_from_stripe(
    user,
    *,
    plan_slug: str | None = None,
    stripe_subscription_id: str = "",
    status: str = UserSubscription.Status.ACTIVE,
    cancel_at_period_end: bool | None = None,
    current_period_end=None,
    keep_existing_plan_if_unknown: bool = False,
) -> UserSubscription:
    """Assign a catalog plan from a Stripe Checkout / subscription event."""
    seed_subscription_plans(update_existing=False)
    existing = (
        UserSubscription.objects.filter(user=user).select_related("plan").first()
    )
    prev_cancel = bool(existing.cancel_at_period_end) if existing else False
    prev_period_end = existing.current_period_end if existing else None
    had_active = bool(
        existing
        and existing.status == UserSubscription.Status.ACTIVE
        and existing.plan_id
    )

    slug = (plan_slug or "").strip().lower()
    plan = get_plan(slug) if slug else None
    if plan is None and keep_existing_plan_if_unknown and existing and existing.plan_id:
        plan = existing.plan
    if plan is None:
        if slug:
            logger.warning(
                "Unknown plan_slug=%s; falling back to Founders for user=%s",
                slug,
                getattr(user, "pk", None),
            )
        plan = get_plan(SubscriptionPlan.Slug.FOUNDERS)
    if plan is None:
        raise RuntimeError("Founders plan missing from catalog")

    # Single auth-event log after plan + cancel fields are applied.
    sub = assign_plan(
        user,
        plan=plan,
        source=UserSubscription.Source.STRIPE,
        status=status,
        stripe_subscription_id=stripe_subscription_id or "",
        log_auth_event=False,
    )
    update_fields: list[str] = []
    if cancel_at_period_end is not None:
        sub.cancel_at_period_end = bool(cancel_at_period_end)
        update_fields.append("cancel_at_period_end")
    if current_period_end is not None:
        sub.current_period_end = current_period_end
        update_fields.append("current_period_end")
    if update_fields:
        sub.save(update_fields=update_fields)

    became_active = (
        sub.status == UserSubscription.Status.ACTIVE and sub.plan_id is not None
    )
    cancel_changed = (
        cancel_at_period_end is not None
        and bool(cancel_at_period_end) != prev_cancel
    )
    period_changed = (
        current_period_end is not None and current_period_end != prev_period_end
    )
    plan_or_status_changed = (
        not existing
        or existing.plan_id != sub.plan_id
        or existing.status != sub.status
        or (existing.source != sub.source)
        or (
            bool(stripe_subscription_id)
            and (existing.stripe_subscription_id or "")
            != (sub.stripe_subscription_id or "")
        )
    )
    if became_active and not had_active:
        log_subscription_auth_event(
            user,
            started=True,
            detail=(
                f"plan={sub.plan.slug} source={sub.source} status={sub.status}"
                f" cancel_at_period_end={sub.cancel_at_period_end}"
            ),
        )
    elif plan_or_status_changed or cancel_changed or period_changed:
        log_subscription_auth_event(
            user,
            started=False,
            detail=(
                f"plan={sub.plan.slug if sub.plan_id else 'none'} "
                f"source={sub.source} status={sub.status} "
                f"cancel_at_period_end={sub.cancel_at_period_end}"
            ),
        )
    return sub


def resolve_plan_from_stripe_price(price_id: str | None) -> SubscriptionPlan | None:
    """Map a Stripe Price id to a local SubscriptionPlan when configured."""
    if not price_id:
        return None
    return SubscriptionPlan.objects.filter(stripe_price_id=price_id).first()


def plan_to_dict(plan: SubscriptionPlan | None) -> dict[str, Any] | None:
    if plan is None:
        return None
    return {
        "slug": plan.slug,
        "name": plan.name,
        "description": plan.description,
        "price_cents": plan.price_cents,
        "currency": plan.currency,
        "interval": plan.interval,
        "is_public": plan.is_public,
        "is_selectable": plan.is_selectable,
        "features": {
            "text_generation": plan.allows_feature("text_generation"),
            "image_generation": plan.allows_feature("image_generation"),
            "rag": plan.allows_feature("rag"),
            "all_future_features": plan.allows_all_future_features,
        },
        "prompt_quota_per_window": plan.prompt_quota_per_window,
        "prompt_window_hours": plan.prompt_window_hours,
        "monthly_token_quota": plan.monthly_token_quota,
        "sort_order": plan.sort_order,
    }