contexta_be/app/routers/billing.py

from fastapi import APIRouter, HTTPException, Depends, Request, Header
from app.models import CheckoutSessionCreate, CheckoutSessionResponse, SubscriptionResponse
from app.database import get_supabase
from app.dependencies import get_current_user
from app.config import settings
from typing import Optional
import logging

logger = logging.getLogger(__name__)
router = APIRouter(prefix="/billing", tags=["Billing"])

PLAN_PRICE_IDS = {
    "starter": settings.stripe_starter_price_id,
    "pro": settings.stripe_pro_price_id,
}


@router.post("/checkout", response_model=CheckoutSessionResponse)
async def create_checkout_session(data: CheckoutSessionCreate, user=Depends(get_current_user)):
    try:
        import stripe
        stripe.api_key = settings.stripe_secret_key

        if data.plan not in PLAN_PRICE_IDS:
            raise HTTPException(status_code=400, detail=f"Invalid plan: {data.plan}")

        price_id = PLAN_PRICE_IDS[data.plan]
        if not price_id:
            raise HTTPException(status_code=400, detail="Plan price not configured")

        supabase = get_supabase()
        sub = supabase.table("subscriptions").select("stripe_customer_id").eq("user_id", user.id).execute()

        customer_id = None
        if sub.data and sub.data[0].get("stripe_customer_id"):
            customer_id = sub.data[0]["stripe_customer_id"]
        else:
            customer = stripe.Customer.create(email=user.email)
            customer_id = customer.id

        session = stripe.checkout.Session.create(
            customer=customer_id,
            payment_method_types=["card"],
            line_items=[{"price": price_id, "quantity": 1}],
            mode="subscription",
            success_url=data.success_url + "?session_id={CHECKOUT_SESSION_ID}",
            cancel_url=data.cancel_url,
            metadata={"user_id": user.id, "plan": data.plan},
        )

        return CheckoutSessionResponse(checkout_url=session.url, session_id=session.id)

    except ImportError:
        raise HTTPException(status_code=500, detail="Stripe not configured")
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Checkout error: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.post("/webhook")
async def stripe_webhook(
    request: Request,
    stripe_signature: Optional[str] = Header(None),
):
    try:
        import stripe

        stripe.api_key = settings.stripe_secret_key
        payload = await request.body()

        if settings.stripe_webhook_secret and stripe_signature:
            try:
                event = stripe.Webhook.construct_event(
                    payload, stripe_signature, settings.stripe_webhook_secret
                )
            except stripe.error.SignatureVerificationError:
                raise HTTPException(status_code=400, detail="Invalid signature")
        else:
            import json
            event = json.loads(payload)

        supabase = get_supabase()
        event_type = event.get("type", "")

        if event_type == "checkout.session.completed":
            session = event["data"]["object"]
            user_id = session.get("metadata", {}).get("user_id")
            plan = session.get("metadata", {}).get("plan", "starter")
            customer_id = session.get("customer")
            subscription_id = session.get("subscription")

            if user_id:
                supabase.table("subscriptions").upsert({
                    "user_id": user_id,
                    "plan": plan,
                    "status": "active",
                    "stripe_customer_id": customer_id,
                    "stripe_subscription_id": subscription_id,
                }, on_conflict="user_id").execute()

        elif event_type in ("customer.subscription.updated", "customer.subscription.deleted"):
            sub_obj = event["data"]["object"]
            customer_id = sub_obj.get("customer")
            status = sub_obj.get("status", "canceled")

            existing = supabase.table("subscriptions").select("*").eq("stripe_customer_id", customer_id).execute()
            if existing.data:
                mapped_status = "active" if status in ("active", "trialing") else "canceled"
                supabase.table("subscriptions").update({
                    "status": mapped_status,
                }).eq("stripe_customer_id", customer_id).execute()

                # Unpublish chatbots if subscription canceled
                if mapped_status == "canceled":
                    user_id = existing.data[0]["user_id"]
                    company = supabase.table("companies").select("id").eq("owner_id", user_id).execute()
                    if company.data:
                        supabase.table("chatbots").update({
                            "is_published": False,
                            "visibility": "preview",
                        }).eq("company_id", company.data[0]["id"]).execute()

        return {"received": True}

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Webhook error: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@router.get("/subscription", response_model=SubscriptionResponse)
async def get_subscription(user=Depends(get_current_user)):
    supabase = get_supabase()
    result = supabase.table("subscriptions").select("*").eq("user_id", user.id).execute()

    if not result.data:
        return SubscriptionResponse(
            id="free",
            user_id=user.id,
            plan="free",
            status="active",
        )

    sub = result.data[0]
    return SubscriptionResponse(
        id=sub["id"],
        user_id=sub["user_id"],
        plan=sub["plan"],
        status=sub["status"],
        stripe_customer_id=sub.get("stripe_customer_id"),
        current_period_start=sub.get("current_period_start"),
        current_period_end=sub.get("current_period_end"),
        chatbots_published=sub.get("chatbots_published", 0),
        conversations_used=sub.get("conversations_used", 0),
        created_at=sub.get("created_at"),
    )


@router.post("/portal")
async def customer_portal(request: Request, user=Depends(get_current_user)):
    """Create Stripe customer portal session"""
    try:
        import stripe
        stripe.api_key = settings.stripe_secret_key

        supabase = get_supabase()
        sub = supabase.table("subscriptions").select("stripe_customer_id").eq("user_id", user.id).execute()

        if not sub.data or not sub.data[0].get("stripe_customer_id"):
            raise HTTPException(status_code=404, detail="No subscription found")

        body = await request.json()
        return_url = body.get("return_url", "http://localhost:5173/settings")

        session = stripe.billing_portal.Session.create(
            customer=sub.data[0]["stripe_customer_id"],
            return_url=return_url,
        )
        return {"url": session.url}

    except HTTPException:
        raise
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))