helper/subscriptions.py

import os
import psycopg
from psycopg.rows import dict_row
from supabase import create_client, Client
import asyncio
import json
from pathlib import Path

ADMINS_PATH = Path("assets/admins.json")

try:
    ADMIN_EMAIL_MAP = json.loads(ADMINS_PATH.read_text())
except Exception as e:
    print(f"Failed to load admins.json: {e}")
    ADMIN_EMAIL_MAP = {}

PLAN_ORDER = ["free", "light", "core", "creator", "professional"]
TIER_CONFIG = {
    "free": {
        "name": "Free Tier",
        "url": "",
        "price": "0.00",
        "limits": {
            "cloudChatDaily": 50,
            "imagesDaily": 10,
            "videosDaily": 3,
            "audioWeekly": 1,
        },
    },
    "light": {
        "name": "InferencePort AI Light",
        "url": "https://buy.stripe.com/bJedR93gk4xw7EFeyabbG08",
        "price": "9.99",
        "limits": {
            "cloudChatDaily": None,
            "imagesDaily": 50,
            "videosDaily": 10,
            "audioWeekly": 5,
        },
    },
    "core": {
        "name": "InferencePort AI Core",
        "url": "https://buy.stripe.com/28E4gzg365BA1gh61EbbG09",
        "price": "15.99",
        "limits": {
            "cloudChatDaily": None,
            "imagesDaily": 150,
            "videosDaily": 30,
            "audioWeekly": 25,
        },
    },
    "creator": {
        "name": "InferencePort AI Creator",
        "url": "https://buy.stripe.com/dRmdR98AE5BA8IJ89MbbG0a",
        "price": "29.99",
        "limits": {
            "cloudChatDaily": None,
            "imagesDaily": 300,
            "videosDaily": 50,
            "audioWeekly": 45,
        },
    },
    "professional": {
        "name": "InferencePort AI Professional",
        "url": "https://buy.stripe.com/14AaEX7wA3ts7EF2PsbbG0b",
        "price": "99.99",
        "limits": {
            "cloudChatDaily": None,
            "imagesDaily": None,
            "videosDaily": None,
            "audioWeekly": 75,
        },
    },
}
USAGE_PERIODS = {
    "cloudChatDaily": "daily",
    "imagesDaily": "daily",
    "videosDaily": "daily",
    "audioWeekly": "weekly",
}
usage_store = {
    "cloudChatDaily": {},
    "imagesDaily": {},
    "videosDaily": {},
    "audioWeekly": {},
}
usage_locks = {
    "cloudChatDaily": {},
    "imagesDaily": {},
    "videosDaily": {},
    "audioWeekly": {},
}

conn = None
conn_lock = asyncio.Lock()

SUPABASE_URL = os.getenv("SUPABASE_URL")
SUPABASE_KEY = os.getenv("SUPABASE_SERVICE_ROLE_KEY")
supabase: Client = create_client(SUPABASE_URL, SUPABASE_KEY)
POSTGRE_SECRET = os.getenv("POSTGRE_SECRET")

conn = psycopg.connect(POSTGRE_SECRET, row_factory=dict_row, sslmode="verify-full", sslrootcert="prod-ca-2021.crt")


async def get_conn():
    global conn

    if conn is None or conn.closed:
        conn = psycopg.connect(POSTGRE_SECRET, row_factory=dict_row, sslmode="verify-full", sslrootcert="prod-ca-2021.crt")

    return conn


def normalize_plan_key(plan_name: str | None) -> str:
    if not plan_name:
        return "free"
    normalized = "".join(ch for ch in str(plan_name).lower() if ch.isalpha())
    if "professional" in normalized:
        return "professional"
    if "creator" in normalized:
        return "creator"
    if "core" in normalized:
        return "core"
    if "light" in normalized:
        return "light"
    return "free"


async def fetch_subscription(jwt: str):
    auth_res = supabase.auth.get_user(jwt)
    if auth_res.user is None:
        return {"error": "Invalid or expired session"}

    user = auth_res.user
    email = user.email

    auth_res = supabase.auth.get_user(jwt)
    if auth_res.user is None:
        return {"error": "Invalid or expired session"}
    
    user = auth_res.user
    email = user.email.lower()
    
    if email in ADMIN_EMAIL_MAP:
        admin = ADMIN_EMAIL_MAP[email]
        print(f"[ADMIN OVERRIDE] {email} → forcing plan '{admin['plan_key']}'")
    
        subscription_obj = {
            "subscription_id": admin.get("subscription_id", "admin-override"),
            "status": admin.get("status", "active"),
            "current_period_end": None,
            "price_id": None,
            "product_name": admin.get("product_name"),
            "nickname": admin.get("nickname"),
            "plan_key": admin.get("plan_key"),
        }
    
        return {
            "email": email,
            "signed_up": user.created_at.isoformat(),
            "subscription": [subscription_obj],
            "plan_key": admin["plan_key"],
        }

    async with conn_lock:
        connection = await get_conn()

        try:
            with connection.cursor() as cur:
                cur.execute("""
                    with cust as (
                        select id 
                        from stripe.customers 
                        where email = %s
                    ),
                    subs as (
                        select
                            s.id as subscription_id,
                            s.status,
                            s.current_period_end,
                            s.items->'data'->0->'price'->>'id' as price_id
                        from stripe.subscriptions s
                        join cust on s.customer = cust.id
                        where s.status in ('active', 'trialing', 'past_due')
                    )
                    select
                        subs.subscription_id,
                        subs.status,
                        subs.current_period_end,
                        subs.price_id,
                        prices.nickname,
                        prices.product as product_id,
                        products.name as product_name
                    from subs
                    left join stripe.prices prices
                      on prices.id = subs.price_id
                    left join stripe.products products
                      on prices.product = products.id;
                """, (email,))
                rows = cur.fetchall()

        except psycopg.OperationalError:
            connection = psycopg.connect(POSTGRE_SECRET, row_factory=dict_row)
            with connection.cursor() as cur:
                cur.execute("""
                    with cust as (
                        select id 
                        from stripe.customers 
                        where email = %s
                    ),
                    subs as (
                        select
                            s.id as subscription_id,
                            s.status,
                            s.current_period_end,
                            s.items->'data'->0->'price'->>'id' as price_id
                        from stripe.subscriptions s
                        join cust on s.customer = cust.id
                        where s.status in ('active', 'trialing', 'past_due')
                    )
                    select
                        subs.subscription_id,
                        subs.status,
                        subs.current_period_end,
                        subs.price_id,
                        prices.nickname,
                        prices.product as product_id,
                        products.name as product_name
                    from subs
                    left join stripe.prices prices
                      on prices.id = subs.price_id
                    left join stripe.products products
                      on prices.product = products.id;
                """, (email,))
                rows = cur.fetchall()

    if not rows:
        return {
            "email": email,
            "signed_up": user.created_at.isoformat(),
            "subscription": None,
            "plan_key": "free",
        }

    subscriptions = []
    preferred_plan_key = "free"

    for row in rows:
        plan_key = normalize_plan_key(row["product_name"] or row["nickname"])

        subscriptions.append({
            "subscription_id": row["subscription_id"],
            "status": row["status"],
            "current_period_end": row["current_period_end"],
            "price_id": row["price_id"],
            "product_name": row["product_name"],
            "nickname": row["nickname"],
            "plan_key": plan_key,
        })

        if row["status"] in ("active", "trialing"):
            preferred_plan_key = plan_key

    return {
        "email": email,
        "signed_up": user.created_at.isoformat(),
        "subscription": subscriptions,
        "plan_key": preferred_plan_key,
    }