Update main.py

main.py

@@ -19,7 +19,7 @@ from elevenlabs import ElevenLabs
 # Import and configure Google GenAI, matching the Streamlit app
 from google import genai
 from google.genai import types
-
+import stripe
 # -----------------------------------------------------------------------------
 # 1. CONFIGURATION & INITIALIZATION
 # -----------------------------------------------------------------------------
@@ -74,7 +74,29 @@ GENERATION_MODEL = "gemini-2.0-flash-exp-image-generation"
 #GENERATION_MODEL = "gemini-2.5-flash-image-preview"
 #TTS_MODEL = "gemini-2.5-flash-preview-tts"
 
-
+# Stripe
+# --- Stripe Initialization ---
+STRIPE_SECRET_KEY = os.environ.get("STRIPE_SECRET_KEY")
+STRIPE_WEBHOOK_SECRET = os.environ.get("STRIPE_WEBHOOK_SECRET")
+
+# Price IDs from Stripe dashboard (price_xxx...) – set these in your env
+STRIPE_PRICE_FIXER = os.environ.get("STRIPE_PRICE_FIXER")  # "The Fixer" (standard)
+STRIPE_PRICE_PRO = os.environ.get("STRIPE_PRICE_PRO")      # "The Pro" (premium)
+
+# Frontend URLs for redirect after Checkout
+STRIPE_SUCCESS_URL = os.environ.get(
+    "STRIPE_SUCCESS_URL",
+    "https://sozofix.tech/billing/success"
+)
+STRIPE_CANCEL_URL = os.environ.get(
+    "STRIPE_CANCEL_URL",
+    "https://sozofix.tech/billing/cancel"
+)
+
+if STRIPE_SECRET_KEY:
+    stripe.api_key = STRIPE_SECRET_KEY
+else:
+    print("WARNING: STRIPE_SECRET_KEY is not set – Stripe endpoints will fail.")
 # -----------------------------------------------------------------------------
 # 2. HELPER FUNCTIONS (Adapted directly from Streamlit App & Template)
 # -----------------------------------------------------------------------------
@@ -233,6 +255,61 @@ import logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+## Stripe
+PLAN_CONFIG = {
+    # The Fixer – standard
+    "standard": {"price_id": STRIPE_PRICE_FIXER, "credits": 200},
+    "fixer":   {"price_id": STRIPE_PRICE_FIXER, "credits": 200},
+
+    # The Pro – premium
+    "premium": {"price_id": STRIPE_PRICE_PRO, "credits": 500},
+    "pro":     {"price_id": STRIPE_PRICE_PRO, "credits": 500},
+}
+
+def get_or_create_stripe_customer(uid: str) -> str:
+    """
+    Ensure the Firebase user has a Stripe customer.
+    Stores the customer id on /users/{uid}/stripeCustomerId.
+    """
+    user_ref = db_ref.child(f'users/{uid}')
+    user_data = user_ref.get() or {}
+
+    existing_id = user_data.get("stripeCustomerId")
+    if existing_id:
+        return existing_id
+
+    email = user_data.get("email")
+    customer = stripe.Customer.create(
+        email=email,
+        metadata={"firebase_uid": uid}
+    )
+
+    user_ref.update({"stripeCustomerId": customer.id})
+    return customer.id
+
+
+def apply_plan_credits(uid: str, plan_key: str):
+    """
+    Adds the correct number of credits for a plan
+    (200 for standard, 500 for premium).
+    """
+    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"]
+    user_ref = db_ref.child(f'users/{uid}')
+    user_data = user_ref.get() or {}
+    current = user_data.get("credits", 0)
+
+    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}"
+    )
+
 # =============================================================================
 # OPEN IMAGE PROXY ENDPOINT (NO AUTHENTICATION)
 # =============================================================================
@@ -1293,8 +1370,109 @@ def log_call_usage(project_id):
         logger.error(f"[LOGGING] A database error occurred for user '{uid}': {e}")
         return jsonify({'error': 'A server error occurred while updating credits.'}), 500
 
+#Stripe Payments
+@app.route("/api/billing/create-checkout-session", methods=["POST"])
+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)
+    """
+    uid = verify_token(request.headers.get("Authorization"))
+    if not uid:
+        return jsonify({"error": "Unauthorized"}), 401
 
+    if not STRIPE_SECRET_KEY or not STRIPE_PRICE_FIXER or not STRIPE_PRICE_PRO:
+        logger.error("[STRIPE] Missing Stripe configuration.")
+        return jsonify({"error": "Stripe is not configured on the server."}), 500
+
+    data = request.get_json() or {}
+    plan = (data.get("plan") or "").lower().strip()
+
+    plan_cfg = PLAN_CONFIG.get(plan)
+    if not plan_cfg or not plan_cfg["price_id"]:
+        return jsonify({"error": "Invalid plan selected."}), 400
+
+    try:
+        customer_id = get_or_create_stripe_customer(uid)
+
+        session = stripe.checkout.Session.create(
+            mode="subscription",
+            customer=customer_id,
+            payment_method_types=["card"],
+            line_items=[{
+                "price": plan_cfg["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,
+            },
+            subscription_data={
+                "metadata": {
+                    "firebase_uid": uid,
+                    "plan": plan,
+                }
+            },
+        )
+
+        logger.info(f"[STRIPE] Created checkout session {session.id} for user {uid}, plan {plan}")
+        return jsonify({
+            "id": session.id,
+            "url": session.url,
+        })
+
+    except Exception as e:
+        logger.error(f"[STRIPE] Error creating checkout session: {e}")
+        return jsonify({"error": "Failed to create checkout session."}), 500
+        
+@app.route("/api/billing/webhook", methods=["POST"])
+def stripe_webhook():
+    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")
+
+    try:
+        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 parsing error: {e}")
+        return "Bad request", 400
+
+    event_type = event.get("type")
+    obj = event.get("data", {}).get("object", {})
+
+    # We only care about successful subscription payments
+    if event_type == "invoice.payment_succeeded":
+        try:
+            subscription_id = obj.get("subscription")
+            if not subscription_id:
+                return "", 200
+
+            sub = stripe.Subscription.retrieve(subscription_id)
+            metadata = sub.get("metadata", {}) or {}
+            uid = metadata.get("firebase_uid")
+            plan = metadata.get("plan")
+
+            if uid and plan:
+                logger.info(f"[STRIPE] invoice.payment_succeeded for user {uid}, plan {plan}")
+                apply_plan_credits(uid, plan)
+        except Exception as e:
+            logger.error(f"[STRIPE] Error handling invoice.payment_succeeded: {e}")
 
+    # You can log other event types here if needed
+    return "", 200
 # -----------------------------------------------------------------------------
 # 7. MAIN EXECUTION
 # -----------------------------------------------------------------------------