scripts/admin_usage_report.py

"""
Knowledge Universe — Admin Usage Report (Upstash REST Version)
Run: python scripts/admin_usage_report.py
"""
import os
import json
from datetime import datetime, timezone
from upstash_redis import Redis
from dotenv import load_dotenv

# Load variables from .env securely
load_dotenv()

# Pull the variables from the environment safely
UPSTASH_URL = os.getenv("UPSTASH_URL")
UPSTASH_TOKEN = os.getenv("UPSTASH_TOKEN")

# Fail-safe: Alert you if the .env variables are missing
if not UPSTASH_URL or not UPSTASH_TOKEN:
    raise ValueError("🚨 ERROR: UPSTASH_URL or UPSTASH_TOKEN is missing from your .env file!")

# Initialize the client securely
client = Redis(url=UPSTASH_URL, token=UPSTASH_TOKEN)

def get_current_month():
    return datetime.now(timezone.utc).strftime("%Y-%m")

def run_report():
    print(f"Connecting to Upstash REST API at {UPSTASH_URL}...")
    
    # 1. Get all customer keys
    customer_keys = client.keys("ku:key:*")
    if not customer_keys:
        print("\nNo customers found in the database yet.")
        return

    users = []
    current_month = get_current_month()

    # 2. Fetch data for each customer
    for key in customer_keys:
        raw_data = client.get(key)
        if not raw_data:
            continue
            
        # Upstash REST client might auto-parse JSON, so we handle both dict and string
        customer = raw_data if isinstance(raw_data, dict) else json.loads(raw_data)
        c_id = customer.get("customer_id")
        
        # Look up their usage for the current month
        usage_key = f"ku:usage:{c_id}:{current_month}"
        usage = client.get(usage_key)
        calls_used = int(usage) if usage else 0
        
        # Format the date
        created_dt = datetime.fromisoformat(customer.get("created_at").replace('Z', '+00:00'))
        date_str = created_dt.strftime("%b %d, %H:%M")

        users.append({
            "email": customer.get("email"),
            "tier": customer.get("tier").upper(),
            "used": calls_used,
            "limit": customer.get("calls_limit"),
            "joined": date_str
        })

    # 3. Sort by usage (highest first) to find the "Whales"
    users.sort(key=lambda x: x["used"], reverse=True)

    # 4. Print the beautiful table
    print(f"\n{'='*80}")
    print(f" KNOWLEDGE UNIVERSE USAGE REPORT ({current_month})")
    print(f"{'='*80}")
    print(f"{'Email':<30} | {'Tier':<8} | {'Used':<6} | {'Limit':<7} | {'Joined'}")
    print("-" * 80)
    
    total_calls = 0
    for u in users:
        total_calls += u['used']
        used_str = f"{u['used']}"
        if u['used'] > 50:
            used_str = f"🔥 {u['used']}"
            
        print(f"{u['email']:<30} | {u['tier']:<8} | {used_str:<6} | {u['limit']:<7} | {u['joined']}")
        
    print("-" * 80)
    print(f"Total Users: {len(users)}  |  Total API Calls This Month: {total_calls}\n")

if __name__ == "__main__":
    run_report()