aza/AzA march 2026/stripe_routes.py

# stripe_routes.py – Stripe billing, webhook handling, license persistence.
#
# IMPORTANT (learned the hard way, 2026-03-27):
#   - Do NOT call .get() on raw StripeObjects; they are not plain dicts.
#   - Do NOT use .to_dict_recursive(); it does not exist in all stripe-python versions.
#   - After Webhook.construct_event(), always work with json.loads(body) instead.
#   - For Subscription.retrieve() / Customer.retrieve(), convert via _stripe_to_dict(obj).
#   - json.loads(str(obj)) breaks when StripeObject contains Decimal values (monetary amounts).
#     Always use _stripe_to_dict() instead. Never json.loads(str(...)) on StripeObjects.
#   - This router is mounted with prefix="/stripe" in main.py.
#     Route decorators here must NOT repeat the prefix (use "/webhook", not "/stripe/webhook").
#   - When changing Stripe env vars on Hetzner, rebuild the container:
#     docker compose up -d --build --force-recreate aza-api   (restart alone is not enough)

from __future__ import annotations

import json
import os
import sqlite3
import time
from dataclasses import dataclass
from decimal import Decimal
from pathlib import Path
from typing import Any, Dict, Optional

import stripe
from fastapi import APIRouter, Header, HTTPException, Request
from fastapi.responses import JSONResponse


STRIPE_SECRET_KEY = os.environ.get("STRIPE_SECRET_KEY", "").strip()
STRIPE_WEBHOOK_SECRET = os.environ.get("STRIPE_WEBHOOK_SECRET", "").strip()
STRIPE_SUCCESS_URL = os.environ.get("STRIPE_SUCCESS_URL", "").strip()
STRIPE_CANCEL_URL = os.environ.get("STRIPE_CANCEL_URL", "").strip()
STRIPE_PORTAL_RETURN_URL = os.environ.get("STRIPE_PORTAL_RETURN_URL", "").strip()

_BASE_DIR = Path(__file__).resolve().parent
DB_PATH = Path(os.environ.get("STRIPE_DB_PATH", str(_BASE_DIR / "data" / "stripe_webhook.sqlite")))
EVENTS_LOG = Path(os.environ.get("STRIPE_EVENTS_LOG", str(_BASE_DIR / "data" / "stripe_events.log.jsonl")))

router = APIRouter(tags=["stripe"])


def _decimal_default(o: Any) -> Any:
    """json.dumps default handler: Decimal → int/float, StripeObject → dict."""
    if isinstance(o, Decimal):
        return int(o) if o == int(o) else float(o)
    if hasattr(o, "to_dict"):
        return o.to_dict()
    if hasattr(o, "keys"):
        return dict(o)
    return str(o)


def _stripe_to_dict(obj: Any) -> dict:
    """Convert a StripeObject to a plain JSON-safe dict.
    Handles Decimal values that break json.loads(str(obj))."""
    raw = dict(obj) if hasattr(obj, "keys") else obj
    return json.loads(json.dumps(raw, default=_decimal_default))


def _require_env() -> None:
    missing = []
    if not STRIPE_SECRET_KEY:
        missing.append("STRIPE_SECRET_KEY")
    if not STRIPE_WEBHOOK_SECRET:
        missing.append("STRIPE_WEBHOOK_SECRET")
    if not STRIPE_SUCCESS_URL:
        missing.append("STRIPE_SUCCESS_URL")
    if not STRIPE_CANCEL_URL:
        missing.append("STRIPE_CANCEL_URL")
    if missing:
        raise HTTPException(status_code=500, detail=f"Stripe misconfigured. Missing env: {', '.join(missing)}")


def _init_stripe() -> None:
    stripe.api_key = STRIPE_SECRET_KEY


def _ensure_storage() -> None:
    DB_PATH.parent.mkdir(parents=True, exist_ok=True)
    EVENTS_LOG.parent.mkdir(parents=True, exist_ok=True)
    with sqlite3.connect(DB_PATH) as con:
        con.execute(
            """
            CREATE TABLE IF NOT EXISTS processed_events (
              event_id TEXT PRIMARY KEY,
              created_at INTEGER NOT NULL
            )
            """
        )
        con.execute(
            """
            CREATE TABLE IF NOT EXISTS licenses (
              subscription_id TEXT PRIMARY KEY,
              customer_id TEXT,
              status TEXT,
              lookup_key TEXT,
              allowed_users INTEGER,
              devices_per_user INTEGER,
              customer_email TEXT,
              client_reference_id TEXT,
              current_period_end INTEGER,
              updated_at INTEGER NOT NULL
            )
            """
        )
        cols = [row[1] for row in con.execute("PRAGMA table_info(licenses)").fetchall()]
        if "current_period_end" not in cols:
            con.execute("ALTER TABLE licenses ADD COLUMN current_period_end INTEGER")
        con.commit()


@dataclass(frozen=True)
class PlanPolicy:
    lookup_key: str
    allowed_users: int
    devices_per_user: int


def _policy_for_lookup_key(lookup_key: str) -> PlanPolicy:
    if lookup_key.startswith("aza_basic_"):
        return PlanPolicy(lookup_key=lookup_key, allowed_users=1, devices_per_user=2)
    if lookup_key.startswith("aza_team_"):
        return PlanPolicy(lookup_key=lookup_key, allowed_users=3, devices_per_user=2)
    return PlanPolicy(lookup_key=lookup_key, allowed_users=1, devices_per_user=1)


_PRICE_TO_LOOKUP: Dict[tuple, str] = {
    ("month", 5900): "aza_basic_monthly",
    ("year", 59000): "aza_basic_yearly",
    ("month", 8900): "aza_team_monthly",
    ("year", 89000): "aza_team_yearly",
}


def _lookup_key_from_price(price: Dict[str, Any]) -> str:
    """Derive lookup_key from price object. Falls back to interval+amount mapping
    when Stripe sandbox omits lookup_key."""
    lk = (price.get("lookup_key") or "").strip()
    if lk:
        return lk
    recurring = price.get("recurring") or {}
    interval = (recurring.get("interval") or "").strip()
    amount = price.get("unit_amount")
    if interval and amount is not None:
        try:
            amount = int(amount)
        except (TypeError, ValueError):
            return ""
        return _PRICE_TO_LOOKUP.get((interval, amount), "")
    return ""


def _price_id_from_lookup_key(lookup_key: str) -> str:
    _init_stripe()
    prices = stripe.Price.list(lookup_keys=[lookup_key], active=True, limit=1)
    if not prices.data:
        raise HTTPException(status_code=400, detail=f"Unknown lookup_key: {lookup_key}")
    return prices.data[0].id


@router.get("/health")
def stripe_health() -> Dict[str, Any]:
    _require_env()
    return {"ok": True}


@router.post("/create_checkout_session")
async def create_checkout_session(payload: Dict[str, Any]) -> Dict[str, Any]:
    _require_env()
    _ensure_storage()
    _init_stripe()

    lookup_key = (payload.get("lookup_key") or "").strip()
    if not lookup_key:
        raise HTTPException(status_code=400, detail="Missing lookup_key")

    customer_email = (payload.get("customer_email") or "").strip() or None
    client_reference_id = (payload.get("client_reference_id") or "").strip() or None

    price_id = _price_id_from_lookup_key(lookup_key)
    policy = _policy_for_lookup_key(lookup_key)

    subscription_data = {
        "metadata": {
            "lookup_key": policy.lookup_key,
            "allowed_users": str(policy.allowed_users),
            "devices_per_user": str(policy.devices_per_user),
        }
    }

    session = stripe.checkout.Session.create(
        mode="subscription",
        line_items=[{"price": price_id, "quantity": 1}],
        success_url=STRIPE_SUCCESS_URL,
        cancel_url=STRIPE_CANCEL_URL,
        customer_email=customer_email,
        client_reference_id=client_reference_id,
        subscription_data=subscription_data,
        allow_promotion_codes=True,
        billing_address_collection="auto",
    )

    return {"id": session.id, "url": session.url}


@router.post("/create_billing_portal_session")
async def create_billing_portal_session(payload: Dict[str, Any]) -> Dict[str, Any]:
    _require_env()
    _init_stripe()
    customer_id = (payload.get("customer_id") or "").strip()
    if not customer_id:
        raise HTTPException(status_code=400, detail="Missing customer_id")
    if not STRIPE_PORTAL_RETURN_URL:
        raise HTTPException(status_code=500, detail="Missing STRIPE_PORTAL_RETURN_URL")

    portal = stripe.billing_portal.Session.create(customer=customer_id, return_url=STRIPE_PORTAL_RETURN_URL)
    return {"url": portal.url}


def _already_processed(event_id: str) -> bool:
    with sqlite3.connect(DB_PATH) as con:
        row = con.execute("SELECT 1 FROM processed_events WHERE event_id=? LIMIT 1", (event_id,)).fetchone()
        return row is not None


def _mark_processed(event_id: str) -> None:
    now = int(time.time())
    with sqlite3.connect(DB_PATH) as con:
        con.execute("INSERT OR IGNORE INTO processed_events(event_id, created_at) VALUES(?, ?)", (event_id, now))
        con.commit()


def _log_event(kind: str, payload: Dict[str, Any]) -> None:
    rec = {"ts": int(time.time()), "kind": kind, "payload": payload}
    with EVENTS_LOG.open("a", encoding="utf-8") as f:
        f.write(json.dumps(rec, ensure_ascii=False, default=_decimal_default) + "\n")


def _upsert_license(
    *,
    subscription_id: str,
    customer_id: str,
    status: str,
    lookup_key: str,
    allowed_users: Optional[int],
    devices_per_user: Optional[int],
    customer_email: Optional[str],
    client_reference_id: Optional[str],
    current_period_end: Optional[int],
) -> None:
    now = int(time.time())
    with sqlite3.connect(DB_PATH) as con:
        con.execute(
            """
            INSERT INTO licenses(
              subscription_id, customer_id, status, lookup_key,
              allowed_users, devices_per_user, customer_email, client_reference_id,
              current_period_end, updated_at
            )
            VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            ON CONFLICT(subscription_id) DO UPDATE SET
              customer_id=excluded.customer_id,
              status=excluded.status,
              lookup_key=excluded.lookup_key,
              allowed_users=excluded.allowed_users,
              devices_per_user=excluded.devices_per_user,
              customer_email=excluded.customer_email,
              client_reference_id=excluded.client_reference_id,
              current_period_end=excluded.current_period_end,
              updated_at=excluded.updated_at
            """,
            (
                subscription_id,
                customer_id,
                status,
                lookup_key,
                allowed_users,
                devices_per_user,
                customer_email,
                client_reference_id,
                current_period_end,
                now,
            ),
        )
        con.commit()


@router.get("/license_debug")
def license_debug(email: str) -> Dict[str, Any]:
    """
    Debug helper (local only): returns latest license row for a given email.
    """
    _ensure_storage()
    email = (email or "").strip().lower()
    if not email:
        raise HTTPException(status_code=400, detail="Missing email")
    with sqlite3.connect(DB_PATH) as con:
        con.row_factory = sqlite3.Row
        row = con.execute(
            """
            SELECT * FROM licenses
            WHERE lower(customer_email)=?
            ORDER BY updated_at DESC
            LIMIT 1
            """,
            (email,),
        ).fetchone()
        if not row:
            raise HTTPException(status_code=404, detail="No license for email")
        return dict(row)


@router.post("/sync_subscription")
async def sync_subscription() -> Dict[str, Any]:
    _require_env()
    _ensure_storage()
    _init_stripe()

    with sqlite3.connect(DB_PATH) as con:
        row = con.execute(
            """
            SELECT subscription_id
            FROM licenses
            WHERE status='active'
            ORDER BY updated_at DESC
            LIMIT 1
            """
        ).fetchone()
        if not row or not row[0]:
            raise HTTPException(status_code=404, detail="No active license found")
        subscription_id = str(row[0])

    sub = _stripe_to_dict(stripe.Subscription.retrieve(subscription_id))
    status = (sub.get("status") or "").strip()
    current_period_end = sub.get("current_period_end")

    if not current_period_end:
        try:
            current_period_end = sub["items"]["data"][0]["current_period_end"]
        except Exception:
            current_period_end = None

    now = int(time.time())
    with sqlite3.connect(DB_PATH) as con:
        con.execute(
            """
            UPDATE licenses
            SET status=?, current_period_end=?, updated_at=?
            WHERE subscription_id=?
            """,
            (status, current_period_end, now, subscription_id),
        )
        con.commit()

    return {
        "subscription_id": subscription_id,
        "status": status,
        "current_period_end": current_period_end,
    }


# Route is "/webhook" here because main.py mounts this router at prefix="/stripe".
@router.post("/webhook")
async def stripe_webhook(
    request: Request,
    stripe_signature: Optional[str] = Header(default=None, alias="Stripe-Signature"),
) -> JSONResponse:
    _require_env()
    _ensure_storage()
    _init_stripe()

    if not stripe_signature:
        raise HTTPException(status_code=400, detail="Missing Stripe-Signature header")

    body = await request.body()
    try:
        stripe.Webhook.construct_event(payload=body, sig_header=stripe_signature, secret=STRIPE_WEBHOOK_SECRET)
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Webhook signature verification failed: {e}")

    # Parse as plain dict – do not use the StripeObject returned by construct_event.
    event = json.loads(body)
    event_id = event.get("id", "")
    etype = event.get("type")

    _log_event("webhook_received", {"event_id": event_id, "type": etype})

    if event_id and _already_processed(event_id):
        _log_event("webhook_duplicate", {"event_id": event_id, "type": etype})
        return JSONResponse({"ok": True, "duplicate": True})

    obj = (event.get("data") or {}).get("object") or {}

    try:
        if etype == "checkout.session.completed":
            subscription_id = obj.get("subscription")
            customer_id = obj.get("customer")
            client_reference_id = obj.get("client_reference_id")

            _log_event("checkout_entered", {
                "event_id": event_id,
                "subscription_id": subscription_id,
                "customer_id": customer_id,
            })

            if not subscription_id or not customer_id:
                _log_event("license_skip", {
                    "event_id": event_id,
                    "reason": "missing_subscription_or_customer",
                    "subscription_id": subscription_id,
                    "customer_id": customer_id,
                })
            else:
                # customer_email: cascade through all known locations.
                customer_email = (
                    obj.get("customer_email")
                    or (obj.get("customer_details") or {}).get("email")
                )
                if not customer_email and customer_id:
                    try:
                        cust = _stripe_to_dict(stripe.Customer.retrieve(customer_id))
                        customer_email = cust.get("email")
                    except Exception:
                        pass
                customer_email = (customer_email or "").strip() or None

                sub = _stripe_to_dict(stripe.Subscription.retrieve(subscription_id, expand=["items.data.price"]))
                status = sub.get("status", "") or ""
                current_period_end = sub.get("current_period_end")
                md = sub.get("metadata") or {}
                lookup_key = (md.get("lookup_key") or "").strip()
                allowed_users = md.get("allowed_users")
                devices_per_user = md.get("devices_per_user")

                if not lookup_key:
                    try:
                        price = sub["items"]["data"][0]["price"]
                        lookup_key = _lookup_key_from_price(price)
                    except Exception:
                        lookup_key = ""

                def _to_int(x: Any) -> Optional[int]:
                    try:
                        return int(x)
                    except Exception:
                        return None

                if lookup_key and not allowed_users:
                    policy = _policy_for_lookup_key(lookup_key)
                    allowed_users = allowed_users or str(policy.allowed_users)
                    devices_per_user = devices_per_user or str(policy.devices_per_user)

                if not lookup_key:
                    _log_event("license_skip", {
                        "event_id": event_id,
                        "reason": "missing_lookup_key",
                        "subscription_id": subscription_id,
                    })

                _upsert_license(
                    subscription_id=subscription_id,
                    customer_id=customer_id,
                    status=status,
                    lookup_key=lookup_key,
                    allowed_users=_to_int(allowed_users),
                    devices_per_user=_to_int(devices_per_user),
                    customer_email=customer_email,
                    client_reference_id=(client_reference_id or "").strip() or None,
                    current_period_end=_to_int(current_period_end),
                )
                _log_event("license_upsert", {
                    "event_id": event_id,
                    "etype": etype,
                    "subscription_id": subscription_id,
                    "status": status,
                    "lookup_key": lookup_key,
                    "email": customer_email,
                    "client_reference_id": client_reference_id,
                })

        elif etype in ("customer.subscription.updated", "customer.subscription.deleted"):
            subscription_id = obj.get("id")
            customer_id = obj.get("customer")
            status = obj.get("status", "") or ("canceled" if etype.endswith("deleted") else "")
            current_period_end = obj.get("current_period_end")
            md = obj.get("metadata") or {}
            lookup_key = (md.get("lookup_key") or "").strip()
            allowed_users = md.get("allowed_users")
            devices_per_user = md.get("devices_per_user")

            if not lookup_key:
                try:
                    items = obj.get("items") or {}
                    data = items.get("data") or []
                    if data:
                        price = data[0].get("price") or {}
                        lookup_key = _lookup_key_from_price(price)
                except Exception:
                    lookup_key = ""

            def _to_int(x: Any) -> Optional[int]:
                try:
                    return int(x)
                except Exception:
                    return None

            if not subscription_id or not customer_id:
                _log_event("license_skip", {
                    "event_id": event_id,
                    "reason": "missing_subscription_or_customer",
                    "etype": etype,
                })
            else:
                _upsert_license(
                    subscription_id=subscription_id,
                    customer_id=customer_id,
                    status=status,
                    lookup_key=lookup_key,
                    allowed_users=_to_int(allowed_users),
                    devices_per_user=_to_int(devices_per_user),
                    customer_email=None,
                    client_reference_id=None,
                    current_period_end=_to_int(current_period_end),
                )
                _log_event("license_upsert", {
                    "event_id": event_id,
                    "etype": etype,
                    "subscription_id": subscription_id,
                    "status": status,
                    "lookup_key": lookup_key,
                })

        else:
            _log_event("stripe_event_ignored", {"event_id": event_id, "type": etype})

        # Only mark as processed AFTER successful handling.
        # Failed events stay unmarked so Stripe can retry them.
        if event_id:
            _mark_processed(event_id)
    except Exception as e:
        _log_event("license_error", {"event_id": event_id, "etype": etype, "error": str(e)})

    return JSONResponse({"ok": True})