from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine, async_sessionmaker from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column from sqlalchemy import BigInteger, String, Float, DateTime, Boolean, Text, Integer, text, select, func, and_ from datetime import datetime, timedelta from typing import Optional, List, Dict, Any import logging logger = logging.getLogger(__name__) class Base(DeclarativeBase): pass class ReferralProgram(Base): __tablename__ = 'referral_programs' id: Mapped[int] = mapped_column(primary_key=True) referrer_id: Mapped[int] = mapped_column(BigInteger, index=True) # Кто пригласил referred_id: Mapped[int] = mapped_column(BigInteger, unique=True, index=True) # Кто был приглашен referral_code: Mapped[str] = mapped_column(String(20), index=True) # Промокод реферера first_reward_paid: Mapped[bool] = mapped_column(Boolean, default=False) # Выплачена ли разовая награда total_earned: Mapped[float] = mapped_column(Float, default=0.0) # Всего заработано created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow) first_reward_at: Mapped[Optional[datetime]] = mapped_column(DateTime) # Когда выплатили первую награду class ReferralEarning(Base): __tablename__ = 'referral_earnings' id: Mapped[int] = mapped_column(primary_key=True) referrer_id: Mapped[int] = mapped_column(BigInteger, index=True) # Кто получил награду referred_id: Mapped[int] = mapped_column(BigInteger, index=True) # От кого получена награда amount: Mapped[float] = mapped_column(Float) # Размер награды earning_type: Mapped[str] = mapped_column(String(20)) # 'first_reward', 'percentage' related_payment_id: Mapped[Optional[int]] = mapped_column(Integer) # Связанный платеж created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow) class User(Base): __tablename__ = 'users' id: Mapped[int] = mapped_column(primary_key=True) telegram_id: Mapped[int] = mapped_column(BigInteger, unique=True, index=True) username: Mapped[Optional[str]] = mapped_column(String(255)) first_name: Mapped[Optional[str]] = mapped_column(String(255)) last_name: Mapped[Optional[str]] = mapped_column(String(255)) language: Mapped[str] = mapped_column(String(10), default='ru') balance: Mapped[float] = mapped_column(Float, default=0.0) is_admin: Mapped[bool] = mapped_column(Boolean, default=False) created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow) remnawave_uuid: Mapped[Optional[str]] = mapped_column(String(255)) is_trial_used: Mapped[bool] = mapped_column(Boolean, default=False) class Subscription(Base): __tablename__ = 'subscriptions' id: Mapped[int] = mapped_column(primary_key=True) name: Mapped[str] = mapped_column(String(255)) description: Mapped[Optional[str]] = mapped_column(Text) price: Mapped[float] = mapped_column(Float) duration_days: Mapped[int] = mapped_column(Integer) traffic_limit_gb: Mapped[int] = mapped_column(Integer, default=0) # 0 = unlimited squad_uuid: Mapped[str] = mapped_column(String(255)) is_active: Mapped[bool] = mapped_column(Boolean, default=True) created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow) is_trial: Mapped[bool] = mapped_column(Boolean, default=False) is_imported: Mapped[bool] = mapped_column(Boolean, default=False) class UserSubscription(Base): __tablename__ = 'user_subscriptions' id: Mapped[int] = mapped_column(primary_key=True) user_id: Mapped[int] = mapped_column(BigInteger, index=True) subscription_id: Mapped[int] = mapped_column(Integer, index=True) short_uuid: Mapped[str] = mapped_column(String(255)) expires_at: Mapped[datetime] = mapped_column(DateTime) is_active: Mapped[bool] = mapped_column(Boolean, default=True) traffic_limit_gb: Mapped[Optional[int]] = mapped_column(Integer) created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow) updated_at: Mapped[Optional[datetime]] = mapped_column(DateTime, onupdate=datetime.utcnow) class Payment(Base): __tablename__ = 'payments' id: Mapped[int] = mapped_column(primary_key=True) user_id: Mapped[int] = mapped_column(BigInteger, index=True) amount: Mapped[float] = mapped_column(Float) payment_type: Mapped[str] = mapped_column(String(50)) description: Mapped[str] = mapped_column(Text) status: Mapped[str] = mapped_column(String(50), default='pending') created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow) class Promocode(Base): __tablename__ = 'promocodes' id: Mapped[int] = mapped_column(primary_key=True) code: Mapped[str] = mapped_column(String(255), unique=True, index=True) discount_amount: Mapped[float] = mapped_column(Float) discount_percent: Mapped[Optional[int]] = mapped_column(Integer) usage_limit: Mapped[int] = mapped_column(Integer, default=1) used_count: Mapped[int] = mapped_column(Integer, default=0) expires_at: Mapped[Optional[datetime]] = mapped_column(DateTime) is_active: Mapped[bool] = mapped_column(Boolean, default=True) created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow) class PromocodeUsage(Base): __tablename__ = 'promocode_usage' id: Mapped[int] = mapped_column(primary_key=True) user_id: Mapped[int] = mapped_column(BigInteger, index=True) promocode_id: Mapped[int] = mapped_column(Integer, index=True) used_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow) class LuckyGame(Base): __tablename__ = 'lucky_games' id: Mapped[int] = mapped_column(primary_key=True) user_id: Mapped[int] = mapped_column(BigInteger, index=True) chosen_number: Mapped[int] = mapped_column(Integer) winning_numbers: Mapped[str] = mapped_column(String(255)) # JSON строка с выигрышными номерами is_winner: Mapped[bool] = mapped_column(Boolean, default=False) reward_amount: Mapped[float] = mapped_column(Float, default=0.0) played_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow) class Database: def __init__(self, database_url: str): self.engine = create_async_engine( database_url, echo=False, pool_pre_ping=True, pool_recycle=300 ) self.session_factory = async_sessionmaker( self.engine, class_=AsyncSession, expire_on_commit=False ) async def init_db(self): async with self.engine.begin() as conn: await conn.run_sync(Base.metadata.create_all) await self.migrate_user_subscriptions() await self.migrate_subscription_imported_field() await self.migrate_referral_tables() async def migrate_referral_tables(self): try: async with self.engine.begin() as conn: await conn.execute(text(""" CREATE TABLE IF NOT EXISTS referral_programs ( id SERIAL PRIMARY KEY, referrer_id BIGINT NOT NULL, referred_id BIGINT UNIQUE NOT NULL, referral_code VARCHAR(20) NOT NULL, first_reward_paid BOOLEAN DEFAULT FALSE, total_earned DOUBLE PRECISION DEFAULT 0.0, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, first_reward_at TIMESTAMP, INDEX idx_referrer (referrer_id), INDEX idx_referred (referred_id), INDEX idx_referral_code (referral_code) ) """)) await conn.execute(text(""" CREATE TABLE IF NOT EXISTS referral_earnings ( id SERIAL PRIMARY KEY, referrer_id BIGINT NOT NULL, referred_id BIGINT NOT NULL, amount DOUBLE PRECISION NOT NULL, earning_type VARCHAR(20) NOT NULL, related_payment_id INTEGER, created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP, INDEX idx_referrer_earnings (referrer_id), INDEX idx_referred_earnings (referred_id), INDEX idx_earning_type (earning_type) ) """)) logger.info("Successfully created referral system tables") except Exception as e: logger.error(f"Error creating referral tables: {e}") async def close(self): await self.engine.dispose() async def get_user_by_telegram_id(self, telegram_id: int) -> Optional[User]: async with self.session_factory() as session: try: from sqlalchemy import select result = await session.execute( select(User).where(User.telegram_id == telegram_id) ) return result.scalar_one_or_none() except Exception as e: logger.error(f"Error getting user by telegram_id {telegram_id}: {e}") return None async def create_user(self, telegram_id: int, username: str = None, first_name: str = None, last_name: str = None, language: str = 'ru', is_admin: bool = False) -> User: async with self.session_factory() as session: try: user = User( telegram_id=telegram_id, username=username, first_name=first_name, last_name=last_name, language=language, is_admin=is_admin ) session.add(user) await session.commit() await session.refresh(user) return user except Exception as e: logger.error(f"Error creating user {telegram_id}: {e}") await session.rollback() raise async def update_user(self, user: User) -> User: async with self.session_factory() as session: try: await session.merge(user) await session.commit() return user except Exception as e: logger.error(f"Error updating user {user.telegram_id}: {e}") await session.rollback() raise async def add_balance(self, user_id: int, amount: float) -> bool: async with self.session_factory() as session: try: from sqlalchemy import select, update result = await session.execute( update(User) .where(User.telegram_id == user_id) .values(balance=User.balance + amount) ) await session.commit() return result.rowcount > 0 except Exception as e: logger.error(f"Error adding balance to user {user_id}: {e}") await session.rollback() return False async def get_all_subscriptions(self, include_inactive: bool = False, exclude_trial: bool = True, exclude_imported: bool = True) -> List[Subscription]: async with self.session_factory() as session: try: from sqlalchemy import select query = select(Subscription) if not include_inactive: query = query.where(Subscription.is_active == True) if exclude_trial: query = query.where(Subscription.is_trial == False) if exclude_imported: query = query.where(Subscription.is_imported == False) result = await session.execute(query) return list(result.scalars().all()) except Exception as e: logger.error(f"Error getting subscriptions: {e}") return [] async def get_all_subscriptions_admin(self) -> List[Subscription]: async with self.session_factory() as session: try: from sqlalchemy import select result = await session.execute(select(Subscription)) return list(result.scalars().all()) except Exception as e: logger.error(f"Error getting admin subscriptions: {e}") return [] async def migrate_subscription_imported_field(self): try: async with self.engine.begin() as conn: try: await conn.execute(text(""" ALTER TABLE subscriptions ADD COLUMN IF NOT EXISTS is_imported BOOLEAN DEFAULT FALSE """)) logger.info("Successfully added is_imported field to subscriptions table") except Exception as e: logger.info(f"Migration may have already been applied: {e}") except Exception as e: logger.error(f"Error during subscription migration: {e}") async def get_subscription_by_id(self, subscription_id: int) -> Optional[Subscription]: async with self.session_factory() as session: try: from sqlalchemy import select result = await session.execute( select(Subscription).where(Subscription.id == subscription_id) ) return result.scalar_one_or_none() except Exception as e: logger.error(f"Error getting subscription {subscription_id}: {e}") return None async def create_subscription(self, name: str, description: str, price: float, duration_days: int, traffic_limit_gb: int, squad_uuid: str, is_imported: bool = False) -> Subscription: async with self.session_factory() as session: try: subscription = Subscription( name=name, description=description, price=price, duration_days=duration_days, traffic_limit_gb=traffic_limit_gb, squad_uuid=squad_uuid, is_imported=is_imported ) session.add(subscription) await session.commit() await session.refresh(subscription) return subscription except Exception as e: logger.error(f"Error creating subscription: {e}") await session.rollback() raise async def update_subscription(self, subscription: Subscription) -> Subscription: async with self.session_factory() as session: try: await session.merge(subscription) await session.commit() return subscription except Exception as e: logger.error(f"Error updating subscription {subscription.id}: {e}") await session.rollback() raise async def delete_subscription(self, subscription_id: int) -> bool: async with self.session_factory() as session: try: from sqlalchemy import delete result = await session.execute( delete(Subscription).where(Subscription.id == subscription_id) ) await session.commit() return result.rowcount > 0 except Exception as e: logger.error(f"Error deleting subscription {subscription_id}: {e}") await session.rollback() return False async def get_user_subscriptions(self, user_id: int) -> List[UserSubscription]: async with self.session_factory() as session: try: from sqlalchemy import select result = await session.execute( select(UserSubscription).where(UserSubscription.user_id == user_id) ) return list(result.scalars().all()) except Exception as e: logger.error(f"Error getting user subscriptions for {user_id}: {e}") return [] async def create_user_subscription(self, user_id: int, subscription_id: int, short_uuid: str, expires_at: datetime, is_active: bool = True, traffic_limit_gb: int = None) -> Optional[UserSubscription]: async with self.session_factory() as session: try: from sqlalchemy import select existing = await session.execute( select(UserSubscription).where( UserSubscription.user_id == user_id, UserSubscription.short_uuid == short_uuid ) ) existing_sub = existing.scalar_one_or_none() if existing_sub: logger.warning(f"Subscription with short_uuid {short_uuid} already exists for user {user_id}") return existing_sub new_subscription = UserSubscription( user_id=user_id, subscription_id=subscription_id, short_uuid=short_uuid, expires_at=expires_at, is_active=is_active ) session.add(new_subscription) await session.commit() await session.refresh(new_subscription) return new_subscription except Exception as e: logger.error(f"Error creating user subscription: {e}") await session.rollback() return None async def create_payment(self, user_id: int, amount: float, payment_type: str, description: str, status: str = 'pending') -> Payment: async with self.session_factory() as session: try: payment = Payment( user_id=user_id, amount=amount, payment_type=payment_type, description=description, status=status ) session.add(payment) await session.commit() await session.refresh(payment) return payment except Exception as e: logger.error(f"Error creating payment: {e}") await session.rollback() raise async def get_payment_by_id(self, payment_id: int) -> Optional[Payment]: async with self.session_factory() as session: try: from sqlalchemy import select result = await session.execute( select(Payment).where(Payment.id == payment_id) ) return result.scalar_one_or_none() except Exception as e: logger.error(f"Error getting payment {payment_id}: {e}") return None async def update_payment(self, payment: Payment) -> Payment: async with self.session_factory() as session: try: await session.merge(payment) await session.commit() return payment except Exception as e: logger.error(f"Error updating payment {payment.id}: {e}") await session.rollback() raise async def get_user_payments(self, user_id: int) -> List[Payment]: async with self.session_factory() as session: try: from sqlalchemy import select, desc result = await session.execute( select(Payment) .where(Payment.user_id == user_id) .order_by(desc(Payment.created_at)) ) return list(result.scalars().all()) except Exception as e: logger.error(f"Error getting user payments for {user_id}: {e}") return [] async def get_promocode_by_code(self, code: str) -> Optional[Promocode]: async with self.session_factory() as session: try: from sqlalchemy import select result = await session.execute( select(Promocode).where(Promocode.code == code) ) return result.scalar_one_or_none() except Exception as e: logger.error(f"Error getting promocode {code}: {e}") return None async def create_promocode(self, code: str, discount_amount: float = 0, discount_percent: int = None, usage_limit: int = 1, expires_at: datetime = None) -> Promocode: async with self.session_factory() as session: try: promocode = Promocode( code=code, discount_amount=discount_amount, discount_percent=discount_percent, usage_limit=usage_limit, expires_at=expires_at ) session.add(promocode) await session.commit() await session.refresh(promocode) return promocode except Exception as e: logger.error(f"Error creating promocode: {e}") await session.rollback() raise async def use_promocode(self, user_id: int, promocode: Promocode) -> bool: async with self.session_factory() as session: try: from sqlalchemy import select existing = await session.execute( select(PromocodeUsage).where( PromocodeUsage.user_id == user_id, PromocodeUsage.promocode_id == promocode.id ) ) if existing.scalar_one_or_none(): return False usage = PromocodeUsage(user_id=user_id, promocode_id=promocode.id) session.add(usage) promocode.used_count += 1 await session.merge(promocode) await session.commit() return True except Exception as e: logger.error(f"Error using promocode: {e}") await session.rollback() return False async def get_all_promocodes(self) -> List[Promocode]: async with self.session_factory() as session: try: from sqlalchemy import select result = await session.execute(select(Promocode)) return list(result.scalars().all()) except Exception as e: logger.error(f"Error getting promocodes: {e}") return [] async def get_all_users(self) -> List[User]: async with self.session_factory() as session: try: from sqlalchemy import select result = await session.execute(select(User)) return list(result.scalars().all()) except Exception as e: logger.error(f"Error getting all users: {e}") return [] async def get_stats(self) -> dict: async with self.session_factory() as session: try: from sqlalchemy import select, func total_users = await session.execute( select(func.count(User.id)) ) total_users = total_users.scalar() total_subs_non_trial = await session.execute( select(func.count(UserSubscription.id)) .join(Subscription, UserSubscription.subscription_id == Subscription.id) .where(Subscription.is_trial == False) ) total_subs_non_trial = total_subs_non_trial.scalar() total_payments = await session.execute( select(func.sum(Payment.amount)).where( Payment.status == 'completed', Payment.payment_type != 'trial' ) ) total_payments = total_payments.scalar() or 0 return { 'total_users': total_users, 'total_subscriptions_non_trial': total_subs_non_trial, 'total_revenue': total_payments } except Exception as e: logger.error(f"Error getting stats: {e}") return { 'total_users': 0, 'total_subscriptions_non_trial': 0, 'total_revenue': 0 } async def get_trial_subscriptions(self) -> List[Subscription]: async with self.session_factory() as session: try: from sqlalchemy import select result = await session.execute( select(Subscription).where(Subscription.is_trial == True) ) return list(result.scalars().all()) except Exception as e: logger.error(f"Error getting trial subscriptions: {e}") return [] async def get_user_subscription_by_short_uuid(self, user_id: int, short_uuid: str) -> Optional[UserSubscription]: async with self.session_factory() as session: try: from sqlalchemy import select result = await session.execute( select(UserSubscription).where( UserSubscription.user_id == user_id, UserSubscription.short_uuid == short_uuid ) ) return result.scalar_one_or_none() except Exception as e: logger.error(f"Error getting user subscription by short_uuid: {e}") return None async def update_user_subscription(self, user_subscription: UserSubscription) -> bool: async with self.session_factory() as session: try: # Устанавливаем время обновления user_subscription.updated_at = datetime.utcnow() # Обновляем подписку await session.merge(user_subscription) await session.commit() return True except Exception as e: logger.error(f"Error updating user subscription: {e}") await session.rollback() return False async def migrate_user_subscriptions(self): try: async with self.engine.begin() as conn: try: await conn.execute(text(""" ALTER TABLE user_subscriptions ADD COLUMN IF NOT EXISTS traffic_limit_gb INTEGER, ADD COLUMN IF NOT EXISTS updated_at TIMESTAMP """)) logger.info("Successfully migrated user_subscriptions table") except Exception as e: logger.info(f"Migration may have already been applied or error occurred: {e}") except Exception as e: logger.error(f"Error during migration: {e}") async def get_expiring_subscriptions(self, user_id: int, days_threshold: int = 3) -> List[UserSubscription]: async with self.session_factory() as session: try: from sqlalchemy import select from datetime import datetime, timedelta threshold_date = datetime.utcnow() + timedelta(days=days_threshold) result = await session.execute( select(UserSubscription).where( UserSubscription.user_id == user_id, UserSubscription.is_active == True, UserSubscription.expires_at <= threshold_date, UserSubscription.expires_at > datetime.utcnow() ) ) return list(result.scalars().all()) except Exception as e: logger.error(f"Error getting expiring subscriptions for {user_id}: {e}") return [] async def has_used_trial(self, user_id: int) -> bool: async with self.session_factory() as session: try: from sqlalchemy import select result = await session.execute( select(User.is_trial_used).where(User.telegram_id == user_id) ) is_trial_used = result.scalar_one_or_none() return is_trial_used or False except Exception as e: logger.error(f"Error checking trial usage for user {user_id}: {e}") return False async def mark_trial_used(self, user_id: int) -> bool: async with self.session_factory() as session: try: from sqlalchemy import update result = await session.execute( update(User) .where(User.telegram_id == user_id) .values(is_trial_used=True) ) await session.commit() return result.rowcount > 0 except Exception as e: logger.error(f"Error marking trial used for user {user_id}: {e}") await session.rollback() return False async def get_all_payments_paginated(self, offset: int = 0, limit: int = 10) -> tuple[List[Payment], int]: async with self.session_factory() as session: try: from sqlalchemy import select, desc, func count_result = await session.execute( select(func.count(Payment.id)) ) total_count = count_result.scalar() result = await session.execute( select(Payment) .order_by(desc(Payment.created_at)) .offset(offset) .limit(limit) ) payments = list(result.scalars().all()) return payments, total_count except Exception as e: logger.error(f"Error getting paginated payments: {e}") return [], 0 async def get_payments_by_type_paginated(self, payment_type: str, offset: int = 0, limit: int = 10) -> tuple[List[Payment], int]: async with self.session_factory() as session: try: from sqlalchemy import select, desc, func count_result = await session.execute( select(func.count(Payment.id)).where(Payment.payment_type == payment_type) ) total_count = count_result.scalar() result = await session.execute( select(Payment) .where(Payment.payment_type == payment_type) .order_by(desc(Payment.created_at)) .offset(offset) .limit(limit) ) payments = list(result.scalars().all()) return payments, total_count except Exception as e: logger.error(f"Error getting paginated payments by type: {e}") return [], 0 async def get_payments_by_status_paginated(self, status: str, offset: int = 0, limit: int = 10) -> tuple[List[Payment], int]: async with self.session_factory() as session: try: from sqlalchemy import select, desc, func count_result = await session.execute( select(func.count(Payment.id)).where(Payment.status == status) ) total_count = count_result.scalar() result = await session.execute( select(Payment) .where(Payment.status == status) .order_by(desc(Payment.created_at)) .offset(offset) .limit(limit) ) payments = list(result.scalars().all()) return payments, total_count except Exception as e: logger.error(f"Error getting paginated payments by status: {e}") return [], 0 async def get_user_subscriptions_by_plan_id(self, plan_id: int) -> List[UserSubscription]: async with self.session_factory() as session: try: from sqlalchemy import select result = await session.execute( select(UserSubscription).where(UserSubscription.subscription_id == plan_id) ) return list(result.scalars().all()) except Exception as e: logger.error(f"Error getting user subscriptions for plan {plan_id}: {e}") return [] async def delete_user_subscription(self, user_subscription_id: int) -> bool: async with self.session_factory() as session: try: from sqlalchemy import delete result = await session.execute( delete(UserSubscription).where(UserSubscription.id == user_subscription_id) ) await session.commit() return result.rowcount > 0 except Exception as e: logger.error(f"Error deleting user subscription {user_subscription_id}: {e}") await session.rollback() return False async def create_referral(self, referrer_id: int, referred_id: int, referral_code: str) -> Optional[ReferralProgram]: async with self.session_factory() as session: try: from sqlalchemy import select if referred_id == 0: placeholder_id = 999999999 - referrer_id existing = await session.execute( select(ReferralProgram).where( ReferralProgram.referrer_id == referrer_id, ReferralProgram.referred_id == placeholder_id ) ) existing_referral = existing.scalar_one_or_none() if existing_referral: logger.info(f"Referral code already exists for user {referrer_id}") return existing_referral referral = ReferralProgram( referrer_id=referrer_id, referred_id=placeholder_id, referral_code=referral_code ) session.add(referral) await session.commit() await session.refresh(referral) logger.info(f"Created referral code storage for user {referrer_id}") return referral existing = await session.execute( select(ReferralProgram).where( ReferralProgram.referred_id == referred_id, ReferralProgram.referred_id < 900000000, ReferralProgram.referred_id > 0 ) ) if existing.scalar_one_or_none(): logger.info(f"User {referred_id} already has a real referrer") return None if referrer_id == referred_id: logger.warning(f"User {referrer_id} tried to refer themselves") return None referral = ReferralProgram( referrer_id=referrer_id, referred_id=referred_id, referral_code=referral_code ) session.add(referral) await session.commit() await session.refresh(referral) logger.info(f"Created real referral: {referrer_id} -> {referred_id}") return referral except Exception as e: logger.error(f"Error creating referral: {e}") await session.rollback() return None async def get_referral_by_referred_id(self, referred_id: int) -> Optional[ReferralProgram]: async with self.session_factory() as session: try: from sqlalchemy import select result = await session.execute( select(ReferralProgram).where(ReferralProgram.referred_id == referred_id) ) return result.scalar_one_or_none() except Exception as e: logger.error(f"Error getting referral: {e}") return None async def get_user_referrals(self, referrer_id: int) -> List[ReferralProgram]: async with self.session_factory() as session: try: from sqlalchemy import select, and_ result = await session.execute( select(ReferralProgram).where( and_( ReferralProgram.referrer_id == referrer_id, ReferralProgram.referred_id < 900000000, ReferralProgram.referred_id > 0 ) ) ) referrals = list(result.scalars().all()) logger.debug(f"Found {len(referrals)} real referrals for user {referrer_id}") return referrals except Exception as e: logger.error(f"Error getting user referrals: {e}") return [] async def create_referral(self, referrer_id: int, referred_id: int, referral_code: str) -> Optional[ReferralProgram]: async with self.session_factory() as session: try: from sqlalchemy import select if referred_id == 0: placeholder_id = 999999999 - referrer_id existing = await session.execute( select(ReferralProgram).where( ReferralProgram.referrer_id == referrer_id, ReferralProgram.referred_id == placeholder_id ) ) existing_referral = existing.scalar_one_or_none() if existing_referral: logger.info(f"Referral code already exists for user {referrer_id}") return existing_referral referral = ReferralProgram( referrer_id=referrer_id, referred_id=placeholder_id, referral_code=referral_code ) session.add(referral) await session.commit() await session.refresh(referral) return referral existing = await session.execute( select(ReferralProgram).where(ReferralProgram.referred_id == referred_id) ) if existing.scalar_one_or_none(): logger.info(f"User {referred_id} already has a referrer") return None if referrer_id == referred_id: logger.warning(f"User {referrer_id} tried to refer themselves") return None referral = ReferralProgram( referrer_id=referrer_id, referred_id=referred_id, referral_code=referral_code ) session.add(referral) await session.commit() await session.refresh(referral) return referral except Exception as e: logger.error(f"Error creating referral: {e}") await session.rollback() return None async def get_user_referral_stats(self, user_id: int) -> Dict: async with self.session_factory() as session: try: from sqlalchemy import select, func, and_, or_ placeholder_id = 999999999 - user_id referrals_count = await session.execute( select(func.count(ReferralProgram.id)) .where( and_( ReferralProgram.referrer_id == user_id, ReferralProgram.referred_id != placeholder_id, ReferralProgram.referred_id != 0 ) ) ) active_referrals = await session.execute( select(func.count(ReferralProgram.id)) .where( and_( ReferralProgram.referrer_id == user_id, ReferralProgram.first_reward_paid == True, ReferralProgram.referred_id != placeholder_id, ReferralProgram.referred_id != 0 ) ) ) total_earned = await session.execute( select(func.sum(ReferralEarning.amount)) .where(ReferralEarning.referrer_id == user_id) ) result = { 'total_referrals': referrals_count.scalar() or 0, 'active_referrals': active_referrals.scalar() or 0, 'total_earned': total_earned.scalar() or 0.0 } logger.info(f"Referral stats for user {user_id}: {result}") return result except Exception as e: logger.error(f"Error getting referral stats: {e}") return { 'total_referrals': 0, 'active_referrals': 0, 'total_earned': 0.0 } async def generate_unique_referral_code(self, user_id: int) -> str: async with self.session_factory() as session: try: import secrets import string base_code = f"REF{user_id}" from sqlalchemy import select existing = await session.execute( select(ReferralProgram).where(ReferralProgram.referral_code == base_code) ) if not existing.scalar_one_or_none(): return base_code for _ in range(10): random_suffix = ''.join(secrets.choice(string.ascii_uppercase + string.digits) for _ in range(4)) code = f"REF{user_id}{random_suffix}" existing = await session.execute( select(ReferralProgram).where(ReferralProgram.referral_code == code) ) if not existing.scalar_one_or_none(): return code return f"REF{''.join(secrets.choice(string.ascii_uppercase + string.digits) for _ in range(8))}" except Exception as e: logger.error(f"Error generating referral code: {e}") return f"REF{user_id}ERR" async def get_user_referrals(self, referrer_id: int) -> List[ReferralProgram]: async with self.session_factory() as session: try: from sqlalchemy import select, and_ placeholder_id = 999999999 - referrer_id result = await session.execute( select(ReferralProgram).where( and_( ReferralProgram.referrer_id == referrer_id, ReferralProgram.referred_id != placeholder_id, ReferralProgram.referred_id != 0 ) ).order_by(ReferralProgram.created_at.desc()) ) referrals = list(result.scalars().all()) logger.info(f"Found {len(referrals)} real referrals for user {referrer_id} (excluding placeholder {placeholder_id})") for ref in referrals: logger.debug(f"Referral: referrer={ref.referrer_id}, referred={ref.referred_id}, " f"first_reward_paid={ref.first_reward_paid}, total_earned={ref.total_earned}") return referrals except Exception as e: logger.error(f"Error getting user referrals: {e}") return [] async def create_referral_earning(self, referrer_id: int, referred_id: int, amount: float, earning_type: str, related_payment_id: Optional[int] = None) -> bool: async with self.session_factory() as session: try: earning = ReferralEarning( referrer_id=referrer_id, referred_id=referred_id, amount=amount, earning_type=earning_type, related_payment_id=related_payment_id ) session.add(earning) from sqlalchemy import select, update referral = await session.execute( select(ReferralProgram).where( ReferralProgram.referrer_id == referrer_id, ReferralProgram.referred_id == referred_id ) ) referral_record = referral.scalar_one_or_none() if referral_record: referral_record.total_earned += amount if earning_type == 'first_reward': referral_record.first_reward_paid = True referral_record.first_reward_at = datetime.utcnow() await session.merge(referral_record) await session.commit() return True except Exception as e: logger.error(f"Error creating referral earning: {e}") await session.rollback() return False async def get_promocode_by_id(self, promocode_id: int) -> Optional[Promocode]: async with self.session_factory() as session: try: from sqlalchemy import select result = await session.execute( select(Promocode).where(Promocode.id == promocode_id) ) return result.scalar_one_or_none() except Exception as e: logger.error(f"Error getting promocode by ID {promocode_id}: {e}") return None async def update_promocode(self, promocode: Promocode) -> Promocode: async with self.session_factory() as session: try: await session.merge(promocode) await session.commit() return promocode except Exception as e: logger.error(f"Error updating promocode {promocode.id}: {e}") await session.rollback() raise async def delete_promocode(self, promocode_id: int) -> bool: async with self.session_factory() as session: try: from sqlalchemy import delete await session.execute( delete(PromocodeUsage).where(PromocodeUsage.promocode_id == promocode_id) ) result = await session.execute( delete(Promocode).where(Promocode.id == promocode_id) ) await session.commit() return result.rowcount > 0 except Exception as e: logger.error(f"Error deleting promocode {promocode_id}: {e}") await session.rollback() return False async def get_regular_promocodes(self, include_inactive: bool = True) -> List[Promocode]: async with self.session_factory() as session: try: from sqlalchemy import select, and_ query = select(Promocode).where(~Promocode.code.startswith('REF')) if not include_inactive: query = query.where(Promocode.is_active == True) result = await session.execute(query.order_by(Promocode.created_at.desc())) return list(result.scalars().all()) except Exception as e: logger.error(f"Error getting regular promocodes: {e}") return [] async def get_expired_promocodes(self) -> List[Promocode]: async with self.session_factory() as session: try: from sqlalchemy import select, and_ from datetime import datetime result = await session.execute( select(Promocode).where( and_( Promocode.expires_at < datetime.utcnow(), ~Promocode.code.startswith('REF') ) ) ) return list(result.scalars().all()) except Exception as e: logger.error(f"Error getting expired promocodes: {e}") return [] async def cleanup_expired_promocodes(self) -> int: async with self.session_factory() as session: try: from sqlalchemy import delete, and_ from datetime import datetime expired_promos = await session.execute( select(Promocode.id).where( and_( Promocode.expires_at < datetime.utcnow(), ~Promocode.code.startswith('REF') ) ) ) expired_ids = [row[0] for row in expired_promos.fetchall()] if not expired_ids: return 0 await session.execute( delete(PromocodeUsage).where(PromocodeUsage.promocode_id.in_(expired_ids)) ) result = await session.execute( delete(Promocode).where(Promocode.id.in_(expired_ids)) ) await session.commit() return result.rowcount except Exception as e: logger.error(f"Error cleaning up expired promocodes: {e}") await session.rollback() return 0 async def deactivate_all_regular_promocodes(self) -> int: async with self.session_factory() as session: try: from sqlalchemy import update, and_ result = await session.execute( update(Promocode) .where( and_( ~Promocode.code.startswith('REF'), # Исключаем реферальные Promocode.is_active == True ) ) .values(is_active=False) ) await session.commit() return result.rowcount except Exception as e: logger.error(f"Error deactivating all promocodes: {e}") await session.rollback() return 0 async def get_promocode_stats(self) -> Dict: async with self.session_factory() as session: try: from sqlalchemy import select, func, and_ from datetime import datetime total_promos = await session.execute( select(func.count(Promocode.id)).where(~Promocode.code.startswith('REF')) ) total_count = total_promos.scalar() or 0 active_promos = await session.execute( select(func.count(Promocode.id)).where( and_( ~Promocode.code.startswith('REF'), Promocode.is_active == True ) ) ) active_count = active_promos.scalar() or 0 expired_promos = await session.execute( select(func.count(Promocode.id)).where( and_( ~Promocode.code.startswith('REF'), Promocode.expires_at < datetime.utcnow() ) ) ) expired_count = expired_promos.scalar() or 0 total_usage = await session.execute( select(func.sum(Promocode.used_count)).where(~Promocode.code.startswith('REF')) ) usage_count = total_usage.scalar() or 0 total_discount = await session.execute( select(func.sum(Promocode.discount_amount * Promocode.used_count)).where( ~Promocode.code.startswith('REF') ) ) discount_amount = total_discount.scalar() or 0.0 top_promos = await session.execute( select(Promocode.code, Promocode.used_count, Promocode.discount_amount) .where(~Promocode.code.startswith('REF')) .order_by(Promocode.used_count.desc()) .limit(5) ) top_promocodes = list(top_promos.fetchall()) return { 'total_promocodes': total_count, 'active_promocodes': active_count, 'expired_promocodes': expired_count, 'total_usage': usage_count, 'total_discount_amount': discount_amount, 'top_promocodes': top_promocodes } except Exception as e: logger.error(f"Error getting promocode stats: {e}") return { 'total_promocodes': 0, 'active_promocodes': 0, 'expired_promocodes': 0, 'total_usage': 0, 'total_discount_amount': 0.0, 'top_promocodes': [] } async def get_promocode_usage_by_id(self, promocode_id: int) -> List[PromocodeUsage]: async with self.session_factory() as session: try: from sqlalchemy import select result = await session.execute( select(PromocodeUsage) .where(PromocodeUsage.promocode_id == promocode_id) .order_by(PromocodeUsage.used_at.desc()) ) return list(result.scalars().all()) except Exception as e: logger.error(f"Error getting promocode usage for {promocode_id}: {e}") return [] async def create_lucky_game(self, user_id: int, chosen_number: int, winning_numbers: List[int], is_winner: bool, reward_amount: float = 0.0) -> Optional[LuckyGame]: async with self.session_factory() as session: try: import json game = LuckyGame( user_id=user_id, chosen_number=chosen_number, winning_numbers=json.dumps(winning_numbers), is_winner=is_winner, reward_amount=reward_amount ) session.add(game) await session.commit() await session.refresh(game) return game except Exception as e: logger.error(f"Error creating lucky game: {e}") await session.rollback() return None async def get_user_last_game_today(self, user_id: int) -> Optional[LuckyGame]: async with self.session_factory() as session: try: from sqlalchemy import select, and_ from datetime import date today = date.today() result = await session.execute( select(LuckyGame) .where( and_( LuckyGame.user_id == user_id, LuckyGame.played_at >= datetime.combine(today, datetime.min.time()), LuckyGame.played_at < datetime.combine(today + timedelta(days=1), datetime.min.time()) ) ) .order_by(LuckyGame.played_at.desc()) .limit(1) ) return result.scalar_one_or_none() except Exception as e: logger.error(f"Error getting user last game today: {e}") return None async def get_user_game_stats(self, user_id: int) -> Dict[str, Any]: async with self.session_factory() as session: try: from sqlalchemy import select, func, and_ # Общее количество игр games_count = await session.execute( select(func.count(LuckyGame.id)) .where(LuckyGame.user_id == user_id) ) total_games = games_count.scalar() or 0 # Количество выигрышей wins_count = await session.execute( select(func.count(LuckyGame.id)) .where( and_( # Вот здесь используется and_ LuckyGame.user_id == user_id, LuckyGame.is_winner == True ) ) ) total_wins = wins_count.scalar() or 0 total_reward = await session.execute( select(func.sum(LuckyGame.reward_amount)) .where(LuckyGame.user_id == user_id) ) total_won = total_reward.scalar() or 0.0 return { 'total_games': total_games, 'total_wins': total_wins, 'total_won': total_won, 'win_rate': (total_wins / total_games * 100) if total_games > 0 else 0 } except Exception as e: logger.error(f"Error getting user game stats: {e}") return { 'total_games': 0, 'total_wins': 0, 'total_won': 0.0, 'win_rate': 0 } async def get_user_game_history(self, user_id: int, limit: int = 10) -> List[Dict[str, Any]]: async with self.session_factory() as session: try: from sqlalchemy import select import json result = await session.execute( select(LuckyGame) .where(LuckyGame.user_id == user_id) .order_by(LuckyGame.played_at.desc()) .limit(limit) ) games = result.scalars().all() history = [] for game in games: winning_numbers = json.loads(game.winning_numbers) if game.winning_numbers else [] history.append({ 'id': game.id, 'chosen_number': game.chosen_number, 'winning_numbers': winning_numbers, 'is_winner': game.is_winner, 'reward_amount': game.reward_amount, 'played_at': game.played_at }) return history except Exception as e: logger.error(f"Error getting user game history: {e}") return [] async def can_play_lucky_game_today(self, user_id: int) -> bool: try: last_game = await self.get_user_last_game_today(user_id) return last_game is None except Exception as e: logger.error(f"Error checking can play today: {e}") return True