database.py

# database.py

import logging
from supabase import create_client, Client
import config

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# --- DATABASE INITIALIZATION ---
try:
    supabase: Client = create_client(config.SUPABASE_URL, config.SUPABASE_KEY)
    logger.info("Successfully connected to Supabase.")
except Exception as e:
    logger.error(f"Error connecting to Supabase: {e}")
    supabase = None

# --- DATABASE FUNCTIONS ---

def get_or_create_user(user_id: int, username: str):
    """
    Retrieves a user from the database or creates a new one if they don't exist.
    """
    if not supabase:
        logger.error("Supabase client not initialized.")
        return None

    try:
        # Check if user exists
        response = supabase.table('users').select('*').eq('user_id', user_id).execute()
        
        # response.data is a list. If it's empty, the user doesn't exist.
        if not response.data:
            logger.info(f"User {user_id} ({username}) not found. Creating new user.")
            insert_response = supabase.table('users').insert({
                'user_id': user_id,
                'username': username
            }).execute()
            
            if not insert_response.data:
                 logger.error(f"Failed to create user {user_id}. Response: {insert_response}")
                 return None
            return insert_response.data[0]
        else:
            logger.info(f"Found existing user {user_id} ({username}).")
            return response.data[0]
            
    except Exception as e:
        logger.error(f"Error in get_or_create_user for user {user_id}: {e}")
        return None

def get_user_credits(user_id: int) -> int:
    """
    Fetches the current credit balance for a given user.
    Returns 0 if the user is not found or an error occurs.
    """
    if not supabase:
        return 0
        
    try:
        response = supabase.table('users').select('credits').eq('user_id', user_id).single().execute()
        if response.data:
            return response.data.get('credits', 0)
        return 0
    except Exception as e:
        logger.error(f"Error getting credits for user {user_id}: {e}")
        return 0

def deduct_credit(user_id: int) -> bool:
    """
    Deducts one credit from a user. Returns True on success, False on failure.
    """
    if not supabase:
        return False

    current_credits = get_user_credits(user_id)
    if current_credits <= 0:
        logger.warning(f"User {user_id} has no credits to deduct.")
        return False

    try:
        new_credits = current_credits - 1
        response = supabase.table('users').update({'credits': new_credits}).eq('user_id', user_id).execute()
        if response.data:
            logger.info(f"Successfully deducted 1 credit from user {user_id}. New balance: {new_credits}")
            return True
        return False
    except Exception as e:
        logger.error(f"Error deducting credit for user {user_id}: {e}")
        return False

def update_user_plan(user_id: int, plan: str, credits_to_add: int):
    """
    Updates a user's plan and adds a specified number of credits to their balance.
    """
    if not supabase:
        return False
        
    try:
        current_credits = get_user_credits(user_id)
        new_credits = current_credits + credits_to_add
        
        response = supabase.table('users').update({
            'plan': plan,
            'credits': new_credits
        }).eq('user_id', user_id).execute()
        
        if response.data:
            logger.info(f"User {user_id} plan updated to '{plan}'. Added {credits_to_add} credits. New balance: {new_credits}")
            return True
        return False
    except Exception as e:
        logger.error(f"Error updating plan for user {user_id}: {e}")
        return False