import datetime import logging from firebase_admin import firestore # We import the manager to get the initialized DB from firebase_manager import firebase_manager logger = logging.getLogger("HamoraServer") class QuotaManagerV2: """ V2 QuotaManager: Uses Firestore exclusively. Ignores local JSON files. """ def __init__(self): self.default_limit = 30 # Can be overridden in config or by user doc self.collection_name = "users" def _get_db(self): return firebase_manager.get_db() def get_today_key(self): return datetime.date.today().isoformat() def get_this_month_key(self): return datetime.date.today().strftime("%Y-%m") def check_limit(self, uid: str, device_id: str = None): """ Returns: (allowed: bool, message: str, current_usage: int, limit: int) """ if not uid: return False, "Missing UID", 0, 0 db = self._get_db() if not db: logger.error("❌ [QUOTA_V2] Firestore DB not available.") return False, "Database error", 0, 0 try: doc_ref = db.collection(self.collection_name).document(uid) doc = doc_ref.get() if not doc.exists: logger.warning(f"⚠️ [QUOTA_V2] User document not found for {uid}. Allowing with default.") return True, "Allowed (Default)", 0, self.default_limit data = doc.to_dict() # Check if blocked if data.get("status") == "rejected": return False, "User is blocked", 0, 0 # Admin or special roles if data.get("role") == "admin" or data.get("status") == "admin": return True, "Unlimited (Admin)", 0, -1 # --- Device ID Enforcement --- if device_id: tier = data.get("tier", "student_basic") allowed_devices = data.get("allowed_devices", []) max_devices = 2 if tier == "parent_premium" else 1 if device_id not in allowed_devices: if len(allowed_devices) >= max_devices: return False, f"Device Limit Exceeded ({max_devices})", 0, 0 else: # Add device atomically doc_ref.update({"allowed_devices": firestore.ArrayUnion([device_id])}) # ----------------------------- limit = data.get("quota_limit", self.default_limit) monthly_token_budget = data.get("monthly_token_budget") # Special limits if limit < 0: return True, "Unlimited", 0, -1 today = self.get_today_key() this_month = self.get_this_month_key() last_usage_date = data.get("last_usage_date", "") last_usage_month = data.get("last_usage_month", "") # Daily Question Quota if last_usage_date != today: current_count = 0 else: current_count = data.get("used_today", 0) if current_count >= limit: return False, f"Daily limit reached ({limit})", current_count, limit # Monthly Token Quota (if defined) if monthly_token_budget is not None: if last_usage_month != this_month: used_tokens = 0 else: used_tokens = data.get("used_tokens_this_month", 0) if used_tokens >= monthly_token_budget: return False, f"Monthly token budget reached", current_count, limit return True, "Allowed", current_count, limit except Exception as e: logger.error(f"❌ [QUOTA_V2] Error checking limit for {uid}: {e}") # Fail closed or open? Let's fail open temporarily so we don't block users if DB hiccups return True, "Error - Fallback Allow", 0, self.default_limit def check_wallet(self, uid: str): """ V5.10.0: Strict token balance check for Pencil Economy (Pre-flight). 1 Pencil = 17,000 tokens. Minimum required for a session: 2,000. """ if not uid or uid == "dev-bypass-user": return True, 1000000 # Unlimited for bypass db = self._get_db() if not db: return False, 0 try: doc_ref = db.collection(self.collection_name).document(uid) doc = doc_ref.get() if not doc.exists: return False, 0 # Must have a document to have a wallet data = doc.to_dict() # V5.10.2: Admin Bypass for Wallet Check if data.get("role") == "admin" or data.get("status") == "admin" or data.get("isAdmin") is True: return True, 1000000 wallet = data.get('wallet', {}) token_balance = wallet.get('token_balance', 0) # ALLOW if balance >= 2000 return token_balance >= 2000, token_balance except Exception as e: logger.error(f"❌ [QUOTA_V2] Failed to check wallet for {uid}: {e}") return True, 0 # Fail open to prevent blocking users on backend hiccups def add_tokens_with_absorption(self, uid: str, tokens_to_add: int): """ V5.10.0: Add tokens with "Debt Absorption" logic. If user is in debt (balance < 0), the debt is cleared and the full amount is added. """ db = self._get_db() if not db: return try: doc_ref = db.collection(self.collection_name).document(uid) @firestore.transactional def update_in_transaction(transaction, doc_ref): snapshot = doc_ref.get(transaction=transaction) current_balance = 0 if snapshot.exists: current_balance = snapshot.to_dict().get('wallet', {}).get('token_balance', 0) # Absorption logic starting_point = max(0, current_balance) new_balance = starting_point + tokens_to_add transaction.update(doc_ref, {"wallet.token_balance": new_balance}) transaction = db.transaction() update_in_transaction(transaction, doc_ref) logger.info(f"💰 [QUOTA_V2] Added {tokens_to_add} tokens to {uid}. Debt absorbed: {current_balance < 0}") except Exception as e: logger.error(f"❌ [QUOTA_V2] Failed to add tokens with absorption for {uid}: {e}") def increment_usage(self, uid: str, increment_questions: int = 1, tokens_used: int = 0): if not uid: return db = self._get_db() if not db: return today = self.get_today_key() this_month = self.get_this_month_key() doc_ref = db.collection(self.collection_name).document(uid) try: @firestore.transactional def update_in_transaction(transaction, doc_ref): snapshot = doc_ref.get(transaction=transaction) if not snapshot.exists: transaction.set(doc_ref, { "quota_limit": self.default_limit, "used_today": increment_questions, "used_tokens_this_month": tokens_used, "last_usage_date": today, "last_usage_month": this_month, "last_seen": firestore.SERVER_TIMESTAMP }) return data = snapshot.to_dict() last_date = data.get("last_usage_date", "") last_month = data.get("last_usage_month", "") # Daily logic if last_date != today: new_usage = increment_questions else: new_usage = data.get("used_today", 0) + increment_questions # Monthly Token logic if last_month != this_month: new_tokens = tokens_used else: new_tokens = data.get("used_tokens_this_month", 0) + tokens_used # V5.9.8: Lifetime Tracking & Economics total_used = data.get("total_tokens_used", 0) + tokens_used # Calculate cost delta from cost_tracker import PRICING # We assume a balanced input/output for a rough estimate if only total_tokens is provided, # or we just use an average price if we don't have the split. # Average price = (0.10 + 0.40) / 2 = 0.25 / 1M tokens avg_price = (PRICING["input"] + PRICING["output"]) / 2 cost_delta = (tokens_used / 1e6) * avg_price total_cost = data.get("total_cost_usd", 0.0) + cost_delta # V5.13.8: Lifetime Tracking (Exercises) total_exercises = data.get("total_exercises_solved", 0) + increment_questions transaction.update(doc_ref, { "used_today": new_usage, "used_tokens_this_month": new_tokens, "total_tokens_used": total_used, "total_exercises_solved": total_exercises, # V5.13.8: Track lifetime solves for Avatar Evolution "total_cost_usd": total_cost, "last_usage_date": today, "last_active_date": today, "last_usage_month": this_month, "last_seen": firestore.SERVER_TIMESTAMP, "wallet.token_balance": firestore.Increment(-tokens_used) }) transaction = db.transaction() update_in_transaction(transaction, doc_ref) logger.info(f"📊 [QUOTA_V2] Incremented usage for {uid}. Date: {today}, Tokens: {tokens_used}") except Exception as e: logger.error(f"❌ [QUOTA_V2] Failed to increment usage for {uid}: {e}") # Global instance quota_manager_v2 = QuotaManagerV2()