Update main.py

main.py

@@ -256,14 +256,31 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 ## Stripe
+# Stripe plan config
 PLAN_CONFIG = {
     # The Fixer – standard
-    "standard": {"price_id": STRIPE_PRICE_FIXER, "credits": 200},
-    "fixer":   {"price_id": STRIPE_PRICE_FIXER, "credits": 200},
+    "standard": {
+        "price_id": STRIPE_PRICE_FIXER,
+        "credits": 200,
+        "label": "The Fixer",
+    },
+    "fixer": {
+        "price_id": STRIPE_PRICE_FIXER,
+        "credits": 200,
+        "label": "The Fixer",
+    },
 
     # The Pro – premium
-    "premium": {"price_id": STRIPE_PRICE_PRO, "credits": 500},
-    "pro":     {"price_id": STRIPE_PRICE_PRO, "credits": 500},
+    "premium": {
+        "price_id": STRIPE_PRICE_PRO,
+        "credits": 500,
+        "label": "The Pro",
+    },
+    "pro": {
+        "price_id": STRIPE_PRICE_PRO,
+        "credits": 500,
+        "label": "The Pro",
+    },
 }
 
 def get_or_create_stripe_customer(uid: str) -> str:
@@ -287,29 +304,86 @@ def get_or_create_stripe_customer(uid: str) -> str:
     user_ref.update({"stripeCustomerId": customer.id})
     return customer.id
 
-
-def apply_plan_credits(uid: str, plan_key: str):
+def reset_plan_credits(uid: str, plan_key: str):
     """
-    Adds the correct number of credits for a plan
-    (200 for standard, 500 for premium).
+    HARD RESET the user's credits to the monthly allowance for the given plan.
+    NO ROLLOVER:
+      - standard / fixer -> 200
+      - premium / pro    -> 500
+
+    Any previous balance is overwritten.
     """
     plan_cfg = PLAN_CONFIG.get(plan_key)
     if not plan_cfg:
         logger.error(f"[STRIPE] Unknown plan '{plan_key}' for user {uid}")
         return
 
-    credits_to_add = plan_cfg["credits"]
+    monthly_credits = float(plan_cfg["credits"])
+
     user_ref = db_ref.child(f'users/{uid}')
     user_data = user_ref.get() or {}
-    current = user_data.get("credits", 0)
+    previous = user_data.get('credits', 0)
+
+    user_ref.update({
+        "credits": monthly_credits,
+        "lastCreditResetAt": datetime.utcnow().isoformat(),
+        "creditResetPlan": plan_key,
+    })
 
-    new_total = current + credits_to_add
-    user_ref.update({"credits": new_total})
     logger.info(
-        f"[STRIPE] Applied {credits_to_add} credits to user {uid} "
-        f"for plan '{plan_key}'. New total: {new_total}"
+        f"[STRIPE] RESET credits for user {uid} to {monthly_credits} "
+        f"for plan '{plan_key}' (no rollover). Previous balance: {previous}"
     )
 
+def update_user_subscription_from_stripe(subscription: dict):
+    """
+    Syncs Stripe subscription info into the Firebase user document.
+
+    Stores:
+      - currentPlan (key: standard/premium)
+      - currentPlanLabel ("The Fixer"/"The Pro")
+      - currentPlanMonthlyCredits (200/500)
+      - planStatus (active/past_due/canceled/unpaid/etc.)
+      - planCurrentPeriodEnd (ISO expiry for current period)
+      - stripeSubscriptionId
+      - planUpdatedAt
+    """
+    metadata = subscription.get("metadata") or {}
+    uid = metadata.get("firebase_uid")
+    plan_key = metadata.get("plan")
+
+    if not uid or not plan_key:
+        logger.error("[STRIPE] Subscription missing firebase_uid or plan in metadata.")
+        return
+
+    plan_cfg = PLAN_CONFIG.get(plan_key, {})
+    plan_label = plan_cfg.get("label", plan_key.title())
+    monthly_credits = plan_cfg.get("credits", 0)
+
+    status = subscription.get("status")  # active, past_due, canceled, unpaid, trialing, etc.
+    current_period_end = subscription.get("current_period_end")  # Unix timestamp
+
+    if current_period_end:
+        expiry_iso = datetime.utcfromtimestamp(current_period_end).isoformat() + "Z"
+    else:
+        expiry_iso = None
+
+    user_ref = db_ref.child(f'users/{uid}')
+    update_data = {
+        "currentPlan": plan_key,                      # "standard" / "premium"
+        "currentPlanLabel": plan_label,               # "The Fixer" / "The Pro"
+        "currentPlanMonthlyCredits": monthly_credits, # 200 / 500
+        "planStatus": status,                         # active / past_due / canceled / unpaid ...
+        "planCurrentPeriodEnd": expiry_iso,           # ISO timestamp of current period end
+        "stripeSubscriptionId": subscription.get("id"),
+        "planUpdatedAt": datetime.utcnow().isoformat(),
+    }
+
+    logger.info(
+        f"[STRIPE] Updating subscription for user {uid}: "
+        f"plan={plan_key}, status={status}, expires={expiry_iso}"
+    )
+    user_ref.update(update_data)
 # =============================================================================
 # OPEN IMAGE PROXY ENDPOINT (NO AUTHENTICATION)
 # =============================================================================
@@ -1375,14 +1449,14 @@ def log_call_usage(project_id):
 def get_stripe_config():
     """
     Returns safe-to-expose Stripe configuration values for the frontend.
-    (No secret keys — only publishable data)
+    No secret keys are exposed here.
     """
     try:
         return jsonify({
             "publishableKey": os.environ.get("STRIPE_PUBLISHABLE_KEY"),
             "priceIds": {
-                "fixer": os.environ.get("STRIPE_PRICE_FIXER"),
-                "pro": os.environ.get("STRIPE_PRICE_PRO"),
+                "fixer": STRIPE_PRICE_FIXER,
+                "pro": STRIPE_PRICE_PRO,
             }
         }), 200
     except Exception as e:
@@ -1394,9 +1468,10 @@ def get_stripe_config():
 def create_checkout_session():
     """
     Creates a Stripe Checkout Session for a recurring subscription.
-    Plans:
-      - 'standard' / 'fixer'  → The Fixer (200 credits / month)
-      - 'premium' / 'pro'     → The Pro (500 credits / month)
+
+    Client body:
+      { "plan": "standard" }  -> The Fixer (200 credits/month)
+      { "plan": "premium" }   -> The Pro   (500 credits/month)
     """
     uid = verify_token(request.headers.get("Authorization"))
     if not uid:
@@ -1407,12 +1482,17 @@ def create_checkout_session():
         return jsonify({"error": "Stripe is not configured on the server."}), 500
 
     data = request.get_json() or {}
-    plan = (data.get("plan") or "").lower().strip()
+    plan_key = (data.get("plan") or "").lower().strip()
 
-    plan_cfg = PLAN_CONFIG.get(plan)
-    if not plan_cfg or not plan_cfg["price_id"]:
+    plan_cfg = PLAN_CONFIG.get(plan_key)
+    if not plan_cfg:
         return jsonify({"error": "Invalid plan selected."}), 400
 
+    price_id = plan_cfg.get("price_id")
+    if not price_id or not price_id.startswith("price_"):
+        logger.error(f"[STRIPE] Misconfigured price_id for plan {plan_key}: {price_id}")
+        return jsonify({"error": "Billing configuration error – contact support."}), 500
+
     try:
         customer_id = get_or_create_stripe_customer(uid)
 
@@ -1421,29 +1501,30 @@ def create_checkout_session():
             customer=customer_id,
             payment_method_types=["card"],
             line_items=[{
-                "price": plan_cfg["price_id"],
+                "price": price_id,
                 "quantity": 1,
             }],
             success_url=STRIPE_SUCCESS_URL + "?session_id={CHECKOUT_SESSION_ID}",
             cancel_url=STRIPE_CANCEL_URL,
-            # So we can find the user/plan again from webhooks
             metadata={
                 "firebase_uid": uid,
-                "plan": plan,
+                "plan": plan_key,
             },
             subscription_data={
                 "metadata": {
                     "firebase_uid": uid,
-                    "plan": plan,
+                    "plan": plan_key,
                 }
             },
         )
 
-        logger.info(f"[STRIPE] Created checkout session {session.id} for user {uid}, plan {plan}")
+        logger.info(
+            f"[STRIPE] Created checkout session {session.id} for user {uid}, plan {plan_key}"
+        )
         return jsonify({
             "id": session.id,
             "url": session.url,
-        })
+        }), 200
 
     except Exception as e:
         logger.error(f"[STRIPE] Error creating checkout session: {e}")
@@ -1451,6 +1532,18 @@ def create_checkout_session():
         
 @app.route("/api/billing/webhook", methods=["POST"])
 def stripe_webhook():
+    """
+    Handles Stripe webhook events:
+      - checkout.session.completed         -> initial subscription + first credits
+      - invoice.payment_succeeded          -> monthly renewals (reset credits)
+      - invoice.payment_failed             -> mark planStatus as past_due/unpaid
+      - customer.subscription.deleted      -> cancelled subscription
+      - customer.subscription.updated      -> keep status/expiry in sync
+    """
+    if not STRIPE_WEBHOOK_SECRET:
+        logger.error("[STRIPE] STRIPE_WEBHOOK_SECRET not set.")
+        return jsonify({"error": "Webhook secret not configured."}), 500
+
     payload = request.data
     sig_header = request.headers.get("Stripe-Signature")
 
@@ -1458,36 +1551,98 @@ def stripe_webhook():
         event = stripe.Webhook.construct_event(
             payload, sig_header, STRIPE_WEBHOOK_SECRET
         )
+    except stripe.error.SignatureVerificationError as e:
+        logger.error(f"[STRIPE] Webhook signature verification failed: {e}")
+        return "Invalid signature", 400
     except Exception as e:
-        logger.error(f"[STRIPE] Webhook error: {e}")
-        return "Bad", 400
+        logger.error(f"[STRIPE] Webhook parsing error: {e}")
+        return "Bad request", 400
 
     event_type = event.get("type")
-    obj = event.get("data", {}).get("object", {})
+    data_obj = event.get("data", {}).get("object", {})
+
     logger.info(f"[STRIPE] Webhook event: {event_type}")
 
-    # Handle subscription billing
-    if event_type == "invoice.payment_succeeded":
-        invoice = obj
-        subscription_id = invoice.get("subscription")
-        if subscription_id:
-            sub = stripe.Subscription.retrieve(subscription_id)
-            md = sub.get("metadata", {}) or {}
-            uid = md.get("firebase_uid")
-            plan = md.get("plan")
-            logger.info(f"[STRIPE] invoice.payment_succeeded uid={uid}, plan={plan}")
-            if uid and plan:
-                apply_plan_credits(uid, plan)
-
-    # Also handle checkout.session.completed as a fallback
-    if event_type == "checkout.session.completed":
-        session = obj
-        md = session.get("metadata", {}) or {}
-        uid = md.get("firebase_uid")
-        plan = md.get("plan")
-        logger.info(f"[STRIPE] checkout.session.completed uid={uid}, plan={plan}")
-        if uid and plan:
-            apply_plan_credits(uid, plan)
+    try:
+        # 1) First payment via Checkout
+        if event_type == "checkout.session.completed":
+            session = data_obj
+            metadata = session.get("metadata") or {}
+            uid = metadata.get("firebase_uid")
+            plan_key = metadata.get("plan")
+            subscription_id = session.get("subscription")
+
+            logger.info(
+                f"[STRIPE] checkout.session.completed uid={uid}, plan={plan_key}, sub={subscription_id}"
+            )
+
+            if subscription_id:
+                sub = stripe.Subscription.retrieve(subscription_id)
+                update_user_subscription_from_stripe(sub)
+
+                # First period: reset credits to plan allowance (no rollover)
+                if uid and plan_key:
+                    reset_plan_credits(uid, plan_key)
+
+        # 2) Every successful invoice (monthly renewal)
+        elif event_type == "invoice.payment_succeeded":
+            invoice = data_obj
+            subscription_id = invoice.get("subscription")
+
+            if subscription_id:
+                sub = stripe.Subscription.retrieve(subscription_id)
+                metadata = sub.get("metadata") or {}
+                uid = metadata.get("firebase_uid")
+                plan_key = metadata.get("plan")
+
+                logger.info(
+                    f"[STRIPE] invoice.payment_succeeded uid={uid}, plan={plan_key}, sub={subscription_id}"
+                )
+
+                update_user_subscription_from_stripe(sub)
+
+                # New billing period: hard reset credits to monthly quota
+                if uid and plan_key:
+                    reset_plan_credits(uid, plan_key)
+
+        # 3) Invoice failed (non-payment) – mark plan as not good
+        elif event_type == "invoice.payment_failed":
+            invoice = data_obj
+            subscription_id = invoice.get("subscription")
+
+            if subscription_id:
+                sub = stripe.Subscription.retrieve(subscription_id)
+                logger.info(
+                    f"[STRIPE] invoice.payment_failed subscription={subscription_id}, status={sub.get('status')}"
+                )
+                # Status will now be past_due/unpaid; keep in sync
+                update_user_subscription_from_stripe(sub)
+
+                # Optional: if you want to be harsh, you could also zero credits here.
+                # metadata = sub.get("metadata") or {}
+                # uid = metadata.get("firebase_uid")
+                # if uid:
+                #     db_ref.child(f'users/{uid}').update({"credits": 0.0})
+                #     logger.info(f"[STRIPE] Payment failed – zeroed credits for user {uid}")
+
+        # 4) Subscription deleted (cancelled)
+        elif event_type == "customer.subscription.deleted":
+            sub = data_obj
+            logger.info(
+                f"[STRIPE] customer.subscription.deleted id={sub.get('id')}, status={sub.get('status')}"
+            )
+            update_user_subscription_from_stripe(sub)
+
+        # 5) Subscription updated (status, period, etc.)
+        elif event_type == "customer.subscription.updated":
+            sub = data_obj
+            logger.info(
+                f"[STRIPE] customer.subscription.updated id={sub.get('id')}, status={sub.get('status')}"
+            )
+            update_user_subscription_from_stripe(sub)
+
+    except Exception as e:
+        logger.error(f"[STRIPE] Error handling webhook {event_type}: {e}")
 
     return "", 200
 # -----------------------------------------------------------------------------