"""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 finance.models import BackerEmail, SubscriptionPlan, UserSubscription logger = logging.getLogger(__name__) # 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_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_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 when available." ), "price_cents": 4000, "is_public": False, "is_selectable": False, "allows_text_generation": True, "allows_image_generation": 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, 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_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_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 = "", ) -> UserSubscription: sub = get_or_create_user_subscription(user) sub.plan = plan sub.source = source sub.status = status if stripe_subscription_id: sub.stripe_subscription_id = stripe_subscription_id sub.save() 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) slug = (plan_slug or "").strip().lower() plan = get_plan(slug) if slug else None if plan is None and keep_existing_plan_if_unknown: existing = UserSubscription.objects.filter(user=user).select_related("plan").first() if 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") sub = assign_plan( user, plan=plan, source=UserSubscription.Source.STRIPE, status=status, stripe_subscription_id=stripe_subscription_id or "", ) 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) 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"), "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, }