auth.py

"""
API Key Authentication & Credit Management
==========================================

Simple in-memory implementation for demo purposes.
Replace with database in production.
"""

from fastapi import HTTPException, Header
from typing import Optional

# ============================================
# In-memory storage (replace with database)
# ============================================

# Demo API keys with starting credits
users_db = {
    "demo_api_key_12345": {
        "credits": 10,
        "email": "demo@example.com"
    }
}


# ============================================
# Authentication
# ============================================

async def verify_api_key(x_api_key: Optional[str] = Header(None)) -> str:
    """
    Verify API key from X-API-Key header.
    
    In production, validate against database.
    """
    if not x_api_key:
        raise HTTPException(
            status_code=401,
            detail={
                "error": "missing_api_key",
                "message": "X-API-Key header is required"
            }
        )
    
    # For demo: accept any key and create user if not exists
    if x_api_key not in users_db:
        # Auto-create user with 5 free credits
        users_db[x_api_key] = {
            "credits": 5,
            "email": None
        }
        print(f"🆕 New user created: {x_api_key[:8]}... (5 free credits)")
    
    return x_api_key


# ============================================
# Credit Management
# ============================================

def get_user_credits(api_key: str) -> int:
    """Get current credit balance for user"""
    if api_key not in users_db:
        return 0
    return users_db[api_key]["credits"]


def deduct_credits(api_key: str, amount: int) -> int:
    """
    Deduct credits from user account.
    Returns remaining credits.
    """
    if api_key not in users_db:
        raise HTTPException(status_code=401, detail="Invalid API key")
    
    current = users_db[api_key]["credits"]
    if current < amount:
        raise HTTPException(
            status_code=402,
            detail=f"Insufficient credits. Have: {current}, Need: {amount}"
        )
    
    users_db[api_key]["credits"] = current - amount
    return users_db[api_key]["credits"]


def add_credits(api_key: str, amount: int) -> int:
    """
    Add credits to user account (after payment).
    Returns new balance.
    """
    if api_key not in users_db:
        users_db[api_key] = {"credits": 0, "email": None}
    
    users_db[api_key]["credits"] += amount
    return users_db[api_key]["credits"]