""" Energy Refill System Handles energy refill purchases using diamonds. Supports multiple tiers including unlimited energy for 24 hours. """ from datetime import datetime, timedelta from typing import Dict, Any import logging from database import get_database_connection from .validation import rate_limit # Refill tier configurations REFILL_TIERS = { 'small': {'energy': 20, 'diamonds': 10}, 'medium': {'energy': 50, 'diamonds': 20}, 'full': {'energy': 100, 'diamonds': 35}, 'unlimited_24h': {'energy': 0, 'diamonds': 50} # Special case } @rate_limit(max_calls=10, time_window=60) def purchase_energy_refill(player_id: int, refill_type: str) -> Dict[str, Any]: """ Handle energy refill purchase. Rate limited to 10 purchases per minute to prevent abuse. Args: player_id: The player's ID refill_type: Type of refill ('small', 'medium', 'full', 'unlimited_24h') Returns: dict with 'success', 'message', and 'new_balance' keys """ conn = None cursor = None try: # Validate refill type if refill_type not in REFILL_TIERS: return { 'success': False, 'message': f'Invalid refill type: {refill_type}' } tier = REFILL_TIERS[refill_type] diamond_cost = tier['diamonds'] energy_amount = tier['energy'] # Get database connection conn = get_database_connection() cursor = conn.cursor(dictionary=True) # Get player current state cursor.execute( "SELECT diamonds, energy, max_energy, unlimited_energy_until FROM players WHERE id = %s", (player_id,) ) player = cursor.fetchone() if not player: return { 'success': False, 'message': 'Player not found' } # Check diamond balance if player['diamonds'] < diamond_cost: return { 'success': False, 'message': f'Not enough diamonds. Need {diamond_cost}, have {player["diamonds"]}' } # Calculate new values unlimited_until = None if refill_type == 'unlimited_24h': # Special case: unlimited energy for 24 hours unlimited_until = datetime.now() + timedelta(hours=24) new_energy = player['max_energy'] # Fill to max immediately # Atomic UPDATE with balance check (concurrent purchase protection) cursor.execute( """UPDATE players SET diamonds = diamonds - %s, energy = %s, unlimited_energy_until = %s WHERE id = %s AND diamonds >= %s""", (diamond_cost, new_energy, unlimited_until, player_id, diamond_cost) ) else: # Normal refill: add energy up to max new_energy = min(player['energy'] + energy_amount, player['max_energy']) # Atomic UPDATE with balance check (concurrent purchase protection) cursor.execute( """UPDATE players SET diamonds = diamonds - %s, energy = %s WHERE id = %s AND diamonds >= %s""", (diamond_cost, new_energy, player_id, diamond_cost) ) # Check if update succeeded (balance was sufficient) if cursor.rowcount == 0: conn.rollback() return { 'success': False, 'message': 'Insufficient diamonds (concurrent purchase or balance changed)' } # Get new balance after atomic update new_diamonds = player['diamonds'] - diamond_cost # Log transaction cursor.execute( """INSERT INTO energy_refill_purchases (player_id, refill_type, energy_amount, diamond_cost, energy_before, energy_after) VALUES (%s, %s, %s, %s, %s, %s)""", (player_id, refill_type, energy_amount, diamond_cost, player['energy'], new_energy) ) conn.commit() logging.info(f"Player {player_id} purchased {refill_type} refill for {diamond_cost} diamonds") return { 'success': True, 'message': 'Energy refilled successfully', 'new_balance': { 'energy': new_energy, 'diamonds': new_diamonds, 'unlimited_until': unlimited_until.isoformat() if unlimited_until else None } } except Exception as e: if conn: conn.rollback() logging.error(f"Energy refill error for player {player_id}: {e}", exc_info=True) return { 'success': False, 'message': 'Server error. Please try again.' } finally: if cursor: cursor.close() if conn: conn.close() def check_unlimited_energy(player_id: int) -> bool: """ Check if player currently has unlimited energy active. Args: player_id: The player's ID Returns: True if unlimited energy is active, False otherwise """ conn = None cursor = None try: conn = get_database_connection() cursor = conn.cursor(dictionary=True) cursor.execute( "SELECT unlimited_energy_until FROM players WHERE id = %s", (player_id,) ) player = cursor.fetchone() if not player or not player['unlimited_energy_until']: return False if datetime.now() < player['unlimited_energy_until']: return True else: # Expired, clear it cursor.execute( "UPDATE players SET unlimited_energy_until = NULL WHERE id = %s", (player_id,) ) conn.commit() return False except Exception as e: logging.error(f"Error checking unlimited energy for player {player_id}: {e}") return False finally: if cursor: cursor.close() if conn: conn.close() def handle_purchase_energy_refill(player_id: int, message_data: Dict[str, Any], send_to_client): """ WebSocket message handler for 'purchaseEnergyRefill' messages. Args: player_id: The player's ID message_data: Message data containing 'refillType' send_to_client: Function to send message back to client """ refill_type = message_data.get('refillType') if not refill_type: send_to_client(player_id, { 'type': 'error', 'error_code': 'INVALID_REQUEST', 'message': 'Missing refillType' }) return # Apply rate limiting (imported from validation module) from .validation import rate_limit result = purchase_energy_refill(player_id, refill_type) if result['success']: # Send success response with updated balances send_to_client(player_id, { 'type': 'purchaseComplete', 'category': 'energy', 'success': True, 'newBalance': result['new_balance'] }) # Also send full player update send_to_client(player_id, { 'type': 'playerUpdate', 'data': { 'energy': result['new_balance']['energy'], 'diamonds': result['new_balance']['diamonds'], 'unlimited_energy_until': result['new_balance']['unlimited_until'] } }) else: # Send error response send_to_client(player_id, { 'type': 'error', 'error_code': 'PURCHASE_FAILED', 'message': result['message'] }) def get_refill_tiers() -> Dict[str, Dict[str, int]]: """ Get all available refill tiers. Returns: Dictionary of refill tiers with their energy and diamond costs """ return REFILL_TIERS.copy()