import logging import re import secrets from datetime import datetime, timedelta, timezone from typing import Any, Dict, List, Optional from sqlalchemy import and_, delete, func, or_, update from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.future import select from sqlalchemy.orm import selectinload from db.models import Subscription, SubscriptionNotification, User INSTALL_SHARE_TOKEN_BYTES = 16 def _subscription_model_payload(sub_payload: Dict[str, Any]) -> Dict[str, Any]: model_columns = Subscription.__mapper__.columns.keys() filtered_payload = {key: value for key, value in sub_payload.items() if key in model_columns} ignored_keys = sorted(set(sub_payload) - set(filtered_payload)) if ignored_keys: logging.warning("Ignoring unsupported subscription payload keys: %s", ignored_keys) return filtered_payload async def get_active_subscription_by_user_id( session: AsyncSession, user_id: int, panel_user_uuid: Optional[str] = None ) -> Optional[Subscription]: stmt = select(Subscription).where( Subscription.user_id == user_id, Subscription.is_active == True, Subscription.end_date > datetime.now(timezone.utc), ) if panel_user_uuid: stmt = stmt.where(Subscription.panel_user_uuid == panel_user_uuid) stmt = stmt.order_by(Subscription.end_date.desc()).limit(1) result = await session.execute(stmt) return result.scalars().first() async def get_subscription_by_panel_subscription_uuid( session: AsyncSession, panel_sub_uuid: str ) -> Optional[Subscription]: stmt = select(Subscription).where(Subscription.panel_subscription_uuid == panel_sub_uuid) result = await session.execute(stmt) return result.scalar_one_or_none() def normalize_install_share_token(value: Any) -> str: token = str(value or "").strip().lower() if not re.fullmatch(r"[a-f0-9]{32}", token): return "" return token async def get_subscription_by_install_share_token( session: AsyncSession, token: str, ) -> Optional[Subscription]: normalized = normalize_install_share_token(token) if not normalized: return None stmt = select(Subscription).where(Subscription.install_share_token == normalized) result = await session.execute(stmt) return result.scalar_one_or_none() async def ensure_install_share_token( session: AsyncSession, subscription: Subscription, ) -> str: raw_existing = str(getattr(subscription, "install_share_token", "") or "").strip() existing = normalize_install_share_token(raw_existing) if existing: if existing != getattr(subscription, "install_share_token", None): subscription.install_share_token = existing await session.flush() return existing subscription_id = getattr(subscription, "subscription_id", None) for _attempt in range(10): token = secrets.token_hex(INSTALL_SHARE_TOKEN_BYTES) if await get_subscription_by_install_share_token(session, token): continue if subscription_id: result = await session.execute( update(Subscription) .where( Subscription.subscription_id == subscription_id, or_( Subscription.install_share_token.is_(None), Subscription.install_share_token == "", Subscription.install_share_token == raw_existing, ), ) .values(install_share_token=token) ) await session.flush() if result.rowcount: await session.refresh(subscription) return ( normalize_install_share_token( getattr(subscription, "install_share_token", None) ) or token ) await session.refresh(subscription) raw_existing = str(getattr(subscription, "install_share_token", "") or "").strip() existing = normalize_install_share_token(raw_existing) if existing: return existing continue subscription.install_share_token = token await session.flush() await session.refresh(subscription) return token raise RuntimeError("Failed to generate a unique install share token") async def get_active_subscriptions_for_user( session: AsyncSession, user_id: int ) -> List[Subscription]: """Get all active subscriptions for a user.""" stmt = ( select(Subscription) .where(Subscription.user_id == user_id, Subscription.is_active == True) .order_by(Subscription.end_date.desc()) ) result = await session.execute(stmt) return result.scalars().all() async def update_subscription( session: AsyncSession, subscription_id: int, update_data: Dict[str, Any] ) -> Optional[Subscription]: sub = await session.get(Subscription, subscription_id) if sub: for key, value in update_data.items(): setattr(sub, key, value) await session.flush() await session.refresh(sub) return sub async def set_auto_renew( session: AsyncSession, subscription_id: int, enabled: bool ) -> Optional[Subscription]: """Toggle auto_renew_enabled for a subscription.""" return await update_subscription(session, subscription_id, {"auto_renew_enabled": enabled}) async def set_user_subscriptions_cancelled_with_grace( session: AsyncSession, user_id: int, grace_days: int = 1 ) -> int: """Mark all active user subscriptions as cancelled with a short grace period. Sets end_date to now + grace_days, status_from_panel to 'CANCELLED', and skip future notifications to reduce noise after cancellation. Returns number of updated rows. """ from datetime import datetime, timedelta, timezone grace_end = datetime.now(timezone.utc) + timedelta(days=grace_days) stmt = ( update(Subscription) .where(Subscription.user_id == user_id, Subscription.is_active == True) .values( end_date=grace_end, status_from_panel="CANCELLED", skip_notifications=True, ) ) result = await session.execute(stmt) return result.rowcount or 0 async def upsert_subscription(session: AsyncSession, sub_payload: Dict[str, Any]) -> Subscription: panel_sub_uuid = sub_payload.get("panel_subscription_uuid") if not panel_sub_uuid: raise ValueError("panel_subscription_uuid is required for upsert.") existing_sub = await get_subscription_by_panel_subscription_uuid(session, panel_sub_uuid) if existing_sub: logging.info( f"Updating existing subscription {existing_sub.subscription_id} by panel_sub_uuid {panel_sub_uuid}" # noqa: E501 ) for key, value in _subscription_model_payload(sub_payload).items(): setattr(existing_sub, key, value) await session.flush() await session.refresh(existing_sub) return existing_sub else: logging.info(f"Creating new subscription with panel_sub_uuid {panel_sub_uuid}") if sub_payload.get("user_id") is None and "panel_user_uuid" not in sub_payload: raise ValueError("For a new subscription without user_id, panel_user_uuid is required.") if "end_date" not in sub_payload: raise ValueError("Missing 'end_date' for new subscription.") if sub_payload.get("user_id") is not None: from .user_dal import get_user_by_id user = await get_user_by_id(session, sub_payload["user_id"]) if not user: raise ValueError( f"User {sub_payload['user_id']} not found for new subscription with panel_uuid {panel_sub_uuid}." # noqa: E501 ) new_sub = Subscription(**_subscription_model_payload(sub_payload)) session.add(new_sub) await session.flush() await session.refresh(new_sub) return new_sub async def deactivate_other_active_subscriptions( session: AsyncSession, panel_user_uuid: str, current_panel_subscription_uuid: Optional[str] ) -> None: stmt = ( update(Subscription) .where( Subscription.panel_user_uuid == panel_user_uuid, Subscription.is_active == True, ) .values(is_active=False, status_from_panel="INACTIVE_BY_BOT_SYNC") ) if current_panel_subscription_uuid: stmt = stmt.where(Subscription.panel_subscription_uuid != current_panel_subscription_uuid) result = await session.execute(stmt) if result.rowcount > 0: logging.info( f"Deactivated {result.rowcount} other active subscriptions for panel_user_uuid {panel_user_uuid}." # noqa: E501 ) async def deactivate_all_user_subscriptions(session: AsyncSession, user_id: int) -> int: stmt = ( update(Subscription) .where(Subscription.user_id == user_id, Subscription.is_active == True) .values(is_active=False, status_from_panel="INACTIVE_USER_NOT_FOUND") ) result = await session.execute(stmt) if result.rowcount > 0: logging.info( f"Deactivated {result.rowcount} subscriptions for user {user_id} due to missing panel user." # noqa: E501 ) return result.rowcount async def delete_all_user_subscriptions(session: AsyncSession, user_id: int) -> int: """Completely delete all user subscriptions.""" stmt = delete(Subscription).where(Subscription.user_id == user_id) result = await session.execute(stmt) if result.rowcount > 0: logging.info( f"Deleted {result.rowcount} subscription records for user {user_id} for trial reset." ) return result.rowcount async def update_subscription_end_date( session: AsyncSession, subscription_id: int, new_end_date: datetime ) -> Optional[Subscription]: return await update_subscription( session, subscription_id, { "end_date": new_end_date, "last_notification_sent": None, "is_active": True, "status_from_panel": "ACTIVE_EXTENDED_BY_BOT", }, ) async def has_any_subscription_for_user(session: AsyncSession, user_id: int) -> bool: stmt = select(Subscription.subscription_id).where(Subscription.user_id == user_id).limit(1) result = await session.execute(stmt) return result.scalar_one_or_none() is not None async def has_trial_blocking_subscription_for_user(session: AsyncSession, user_id: int) -> bool: now_utc = datetime.now(timezone.utc) reset_at = ( select(User.trial_eligibility_reset_at).where(User.user_id == user_id).scalar_subquery() ) subscription_anchor = func.coalesce(Subscription.start_date, Subscription.end_date) stmt = ( select(Subscription.subscription_id) .where( Subscription.user_id == user_id, or_( reset_at.is_(None), and_(Subscription.is_active == True, Subscription.end_date > now_utc), subscription_anchor > reset_at, ), ) .limit(1) ) result = await session.execute(stmt) return result.scalar_one_or_none() is not None async def get_subscriptions_near_expiration( session: AsyncSession, days_threshold: int ) -> List[Subscription]: now_utc = datetime.now(timezone.utc) threshold_date = now_utc + timedelta(days=days_threshold) stmt = ( select(Subscription) .join(Subscription.user) .where( Subscription.is_active == True, Subscription.skip_notifications == False, Subscription.end_date > now_utc, Subscription.end_date <= threshold_date, or_( Subscription.last_notification_sent == None, func.date(Subscription.last_notification_sent) < func.date(now_utc), ), ) .order_by(Subscription.end_date.asc()) .options(selectinload(Subscription.user)) ) result = await session.execute(stmt) return result.scalars().all() async def update_subscription_notification_time( session: AsyncSession, subscription_id: int, notification_time: datetime ) -> Optional[Subscription]: return await update_subscription( session, subscription_id, {"last_notification_sent": notification_time} ) async def has_subscription_notification( session: AsyncSession, subscription_id: int, notification_key: str, ) -> bool: stmt = ( select(SubscriptionNotification.notification_id) .where( SubscriptionNotification.subscription_id == subscription_id, SubscriptionNotification.notification_key == notification_key, ) .limit(1) ) result = await session.execute(stmt) return result.scalar_one_or_none() is not None async def record_subscription_notification( session: AsyncSession, subscription_id: int, notification_key: str, *, sent_at: Optional[datetime] = None, ) -> None: if sent_at is None: sent_at = datetime.now(timezone.utc) existing = await has_subscription_notification(session, subscription_id, notification_key) if existing: return session.add( SubscriptionNotification( subscription_id=subscription_id, notification_key=notification_key, sent_at=sent_at, ) ) await update_subscription_notification_time(session, subscription_id, sent_at) async def find_subscription_for_notification_update( session: AsyncSession, user_id: int, subscription_end_date_to_match: datetime ) -> Optional[Subscription]: if subscription_end_date_to_match.tzinfo is None: subscription_end_date_to_match = subscription_end_date_to_match.replace(tzinfo=timezone.utc) stmt = ( select(Subscription) .where( Subscription.user_id == user_id, Subscription.is_active == True, Subscription.end_date >= subscription_end_date_to_match - timedelta(seconds=1), Subscription.end_date <= subscription_end_date_to_match + timedelta(seconds=1), ) .limit(1) ) result = await session.execute(stmt) return result.scalar_one_or_none()