import logging import secrets import string from typing import Optional, List, Dict, Any, Tuple from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.future import select from sqlalchemy.orm import selectinload from sqlalchemy import update, delete, func, and_, or_, desc from sqlalchemy.orm import aliased from datetime import datetime, timezone from sqlalchemy.dialects.postgresql import insert as pg_insert from ..models import ( User, Subscription, Payment, PromoCodeActivation, MessageLog, UserBilling, UserPaymentMethod, AdAttribution, ) REFERRAL_CODE_ALPHABET = string.ascii_uppercase + string.digits REFERRAL_CODE_LENGTH = 9 MAX_REFERRAL_CODE_ATTEMPTS = 25 def _generate_referral_code_candidate() -> str: return "".join( secrets.choice(REFERRAL_CODE_ALPHABET) for _ in range(REFERRAL_CODE_LENGTH) ) async def _referral_code_exists(session: AsyncSession, code: str) -> bool: stmt = select(User.user_id).where(User.referral_code == code) result = await session.execute(stmt) return result.scalar_one_or_none() is not None async def generate_unique_referral_code(session: AsyncSession) -> str: """ Generate a unique referral code consisting of uppercase alphanumeric characters. Retries until a free code is found or raises RuntimeError after exceeding attempts. """ for _ in range(MAX_REFERRAL_CODE_ATTEMPTS): candidate = _generate_referral_code_candidate() if not await _referral_code_exists(session, candidate): return candidate raise RuntimeError("Failed to generate a unique referral code after several attempts.") async def ensure_referral_code(session: AsyncSession, user: User) -> str: """ Ensure the provided user has a referral code, generating and persisting it if missing. Returns the existing or newly generated code. """ if user.referral_code: normalized = user.referral_code.strip().upper() if normalized != user.referral_code: user.referral_code = normalized await session.flush() await session.refresh(user) return user.referral_code user.referral_code = await generate_unique_referral_code(session) await session.flush() await session.refresh(user) return user.referral_code async def get_user_by_id(session: AsyncSession, user_id: int) -> Optional[User]: stmt = select(User).where(User.user_id == user_id) result = await session.execute(stmt) return result.scalar_one_or_none() async def get_user_by_username(session: AsyncSession, username: str) -> Optional[User]: clean_username = username.lstrip("@").lower() stmt = select(User).where(func.lower(User.username) == clean_username) result = await session.execute(stmt) return result.scalar_one_or_none() async def get_user_by_panel_uuid( session: AsyncSession, panel_uuid: str ) -> Optional[User]: stmt = select(User).where(User.panel_user_uuid == panel_uuid) result = await session.execute(stmt) return result.scalar_one_or_none() ## Removed unused generic get_user helper to keep DAL explicit and simple async def create_user(session: AsyncSession, user_data: Dict[str, Any]) -> Tuple[User, bool]: """Create a user if not exists in a race-safe way. Returns a tuple of (user, created_flag). """ if "registration_date" not in user_data: user_data["registration_date"] = datetime.now(timezone.utc) if not user_data.get("referral_code"): user_data["referral_code"] = await generate_unique_referral_code(session) else: user_data["referral_code"] = user_data["referral_code"].strip().upper() # Use PostgreSQL upsert to avoid IntegrityError on concurrent inserts stmt = ( pg_insert(User) .values(**user_data) .on_conflict_do_nothing(index_elements=[User.user_id]) .returning(User.user_id) ) result = await session.execute(stmt) inserted_row = result.first() created = inserted_row is not None # Fetch the user (inserted just now or pre-existing) user_id: int = user_data["user_id"] user = await get_user_by_id(session, user_id) if created and user is not None: logging.info( f"New user {user.user_id} created in DAL. Referred by: {user.referred_by_id or 'N/A'}." ) elif user is not None: logging.info( f"User {user.user_id} already exists in DAL. Proceeding without creation." ) return user, created async def get_user_by_referral_code(session: AsyncSession, referral_code: str) -> Optional[User]: normalized = referral_code.strip().upper() if not normalized: return None stmt = select(User).where(User.referral_code == normalized) result = await session.execute(stmt) return result.scalar_one_or_none() async def update_user( session: AsyncSession, user_id: int, update_data: Dict[str, Any] ) -> Optional[User]: user = await get_user_by_id(session, user_id) if user: for key, value in update_data.items(): setattr(user, key, value) await session.flush() await session.refresh(user) return user async def update_user_language( session: AsyncSession, user_id: int, lang_code: str ) -> bool: stmt = update(User).where(User.user_id == user_id).values(language_code=lang_code) result = await session.execute(stmt) return result.rowcount > 0 async def get_banned_users(session: AsyncSession) -> List[User]: """Get all banned users""" stmt = ( select(User) .where(User.is_banned == True) .order_by(User.registration_date.desc()) ) result = await session.execute(stmt) return result.scalars().all() async def get_all_users_paginated( session: AsyncSession, *, page: int = 0, page_size: int = 15 ) -> List[User]: """Return a slice of users ordered by newest registration first.""" safe_page = max(page, 0) safe_page_size = max(page_size, 1) stmt = ( select(User) .order_by(User.registration_date.desc()) .offset(safe_page * safe_page_size) .limit(safe_page_size) ) result = await session.execute(stmt) return result.scalars().all() async def count_all_users(session: AsyncSession) -> int: """Count total number of users.""" result = await session.execute(select(func.count(User.user_id))) return result.scalar_one() async def get_all_active_user_ids_for_broadcast(session: AsyncSession) -> List[int]: stmt = select(User.user_id).where(User.is_banned == False) result = await session.execute(stmt) return result.scalars().all() async def get_all_users_with_panel_uuid(session: AsyncSession) -> List[User]: stmt = select(User).where(User.panel_user_uuid.is_not(None)) result = await session.execute(stmt) return result.scalars().all() async def get_enhanced_user_statistics(session: AsyncSession) -> Dict[str, Any]: """Get comprehensive user statistics including active users, trial users, etc.""" from datetime import datetime, timezone # Use timezone-aware UTC to avoid naive/aware comparison issues in SQL queries now = datetime.now(timezone.utc) today_start = now.replace(hour=0, minute=0, second=0, microsecond=0) # Total users total_users_stmt = select(func.count(User.user_id)) total_users = (await session.execute(total_users_stmt)).scalar() or 0 # Banned users banned_users_stmt = select(func.count(User.user_id)).where(User.is_banned == True) banned_users = (await session.execute(banned_users_stmt)).scalar() or 0 # Active users today (proxy: registered today) active_today_stmt = select(func.count(User.user_id)).where(User.registration_date >= today_start) active_today = (await session.execute(active_today_stmt)).scalar() or 0 # Users with active paid subscriptions (non-trial providers only) paid_subs_stmt = ( select(func.count(func.distinct(Subscription.user_id))) .join(User, Subscription.user_id == User.user_id) .where( and_( Subscription.is_active == True, Subscription.end_date > now, Subscription.provider.is_not(None) # Not trial ) ) ) paid_subs_users = (await session.execute(paid_subs_stmt)).scalar() or 0 # Users on trial period trial_subs_stmt = ( select(func.count(func.distinct(Subscription.user_id))) .join(User, Subscription.user_id == User.user_id) .where( and_( Subscription.is_active == True, Subscription.end_date > now, Subscription.provider.is_(None) # Trial subscriptions ) ) ) trial_users = (await session.execute(trial_subs_stmt)).scalar() or 0 # Inactive users (no active subscription) inactive_users = total_users - paid_subs_users - trial_users - banned_users # Users attracted via referral referral_users_stmt = select(func.count(User.user_id)).where(User.referred_by_id.is_not(None)) referral_users = (await session.execute(referral_users_stmt)).scalar() or 0 return { "total_users": total_users, "banned_users": banned_users, "active_today": active_today, "paid_subscriptions": paid_subs_users, "trial_users": trial_users, "inactive_users": max(0, inactive_users), "referral_users": referral_users } async def get_user_ids_with_active_subscription(session: AsyncSession) -> List[int]: """Return non-banned user IDs who have an active subscription (paid or trial).""" from datetime import datetime, timezone now = datetime.now(timezone.utc) stmt = ( select(func.distinct(Subscription.user_id)) .join(User, Subscription.user_id == User.user_id) .where( and_( User.is_banned == False, Subscription.is_active == True, Subscription.end_date > now, ) ) ) result = await session.execute(stmt) return result.scalars().all() async def get_user_ids_without_active_subscription(session: AsyncSession) -> List[int]: """Return non-banned user IDs who do NOT have any active subscription.""" from datetime import datetime, timezone now = datetime.now(timezone.utc) # Subquery for users with active subscription active_subs_subq = ( select(Subscription.user_id) .where( and_( Subscription.is_active == True, Subscription.end_date > now, ) ) ).scalar_subquery() stmt = ( select(User.user_id) .where( and_( User.is_banned == False, ~User.user_id.in_(active_subs_subq), ) ) ) result = await session.execute(stmt) return result.scalars().all() async def delete_user_and_relations(session: AsyncSession, user_id: int) -> bool: """Completely remove a user and all dependent records from the database. This helper ensures we do not leave dangling foreign keys or orphaned data. """ user = await get_user_by_id(session, user_id) if not user: return False # Ensure referral pointers do not block deletion await session.execute( update(User).where(User.referred_by_id == user_id).values(referred_by_id=None) ) # Clean up dependent tables that do not cascade automatically await session.execute( delete(MessageLog).where( or_(MessageLog.user_id == user_id, MessageLog.target_user_id == user_id) ) ) await session.execute(delete(Payment).where(Payment.user_id == user_id)) await session.execute( delete(Subscription).where(Subscription.user_id == user_id) ) await session.execute( delete(PromoCodeActivation).where(PromoCodeActivation.user_id == user_id) ) await session.execute( delete(UserPaymentMethod).where(UserPaymentMethod.user_id == user_id) ) await session.execute(delete(UserBilling).where(UserBilling.user_id == user_id)) await session.execute(delete(AdAttribution).where(AdAttribution.user_id == user_id)) await session.delete(user) await session.flush() return True async def get_top_users_by_traffic_used( session: AsyncSession, *, limit: int = 10, ) -> List[Dict[str, Any]]: """Return top users by total used traffic across all subscriptions.""" safe_limit = max(1, limit) total_traffic_used = func.coalesce(func.sum(Subscription.traffic_used_bytes), 0) stmt = ( select( User.user_id, User.username, User.first_name, total_traffic_used.label("traffic_used_bytes"), ) .join(Subscription, Subscription.user_id == User.user_id, isouter=True) .group_by(User.user_id, User.username, User.first_name) .having(total_traffic_used > 0) .order_by(desc("traffic_used_bytes"), User.user_id.asc()) .limit(safe_limit) ) result = await session.execute(stmt) return [dict(row._mapping) for row in result] async def get_top_users_by_lifetime_traffic_used( session: AsyncSession, *, limit: int = 10, ) -> List[Dict[str, Any]]: """Return top users by lifetime used traffic from panel data.""" safe_limit = max(1, limit) lifetime_used = func.coalesce(User.lifetime_used_traffic_bytes, 0) stmt = ( select( User.user_id, User.username, User.first_name, lifetime_used.label("lifetime_used_traffic_bytes"), ) .where(lifetime_used > 0) .order_by(desc("lifetime_used_traffic_bytes"), User.user_id.asc()) .limit(safe_limit) ) result = await session.execute(stmt) return [dict(row._mapping) for row in result] async def get_top_users_by_referrals_count( session: AsyncSession, *, limit: int = 10, ) -> List[Dict[str, Any]]: """Return top users by number of invited users.""" safe_limit = max(1, limit) referred_user = aliased(User) invited_count = func.count(referred_user.user_id) stmt = ( select( User.user_id, User.username, User.first_name, invited_count.label("invited_count"), ) .join(referred_user, referred_user.referred_by_id == User.user_id, isouter=True) .group_by(User.user_id, User.username, User.first_name) .having(invited_count > 0) .order_by(desc("invited_count"), User.user_id.asc()) .limit(safe_limit) ) result = await session.execute(stmt) return [dict(row._mapping) for row in result] async def get_top_users_by_referral_revenue( session: AsyncSession, *, limit: int = 10, ) -> List[Dict[str, Any]]: """Return top users by total revenue brought by all invited users.""" safe_limit = max(1, limit) referred_user = aliased(User) referral_revenue = func.coalesce(func.sum(Payment.amount), 0.0) stmt = ( select( User.user_id, User.username, User.first_name, referral_revenue.label("referral_revenue"), ) .join(referred_user, referred_user.referred_by_id == User.user_id, isouter=True) .join( Payment, and_( Payment.user_id == referred_user.user_id, Payment.status == "succeeded", ), isouter=True, ) .group_by(User.user_id, User.username, User.first_name) .having(referral_revenue > 0) .order_by(desc("referral_revenue"), User.user_id.asc()) .limit(safe_limit) ) result = await session.execute(stmt) return [dict(row._mapping) for row in result]