from datetime import datetime from sqlalchemy import delete, func, select, text, update from sqlalchemy.exc import SQLAlchemyError from sqlalchemy.ext.asyncio import AsyncSession from database.models import Key, User from database.users import invalidate_user_snapshot from logger import logger async def store_key( session: AsyncSession, tg_id: int, client_id: str, email: str, expiry_time: int, key: str, server_id: str, remnawave_link: str = None, tariff_id: int | None = None, alias: str | None = None, selected_device_limit: int | None = None, selected_traffic_limit: int | None = None, selected_price_rub: int | None = None, current_device_limit: int | None = None, current_traffic_limit: int | None = None, ): """Сохраняет или обновляет ключ подписки.""" try: exists = await session.execute(select(Key).where(Key.tg_id == tg_id, Key.client_id == client_id)) existing_key = exists.scalar_one_or_none() if existing_key: values: dict = { "email": email, "expiry_time": expiry_time, "key": key, "server_id": server_id, "remnawave_link": remnawave_link, "tariff_id": tariff_id, "alias": alias, } if selected_device_limit is not None: values["selected_device_limit"] = selected_device_limit if selected_traffic_limit is not None: values["selected_traffic_limit"] = selected_traffic_limit if selected_price_rub is not None: values["selected_price_rub"] = selected_price_rub if current_device_limit is not None: values["current_device_limit"] = current_device_limit if current_traffic_limit is not None: values["current_traffic_limit"] = current_traffic_limit await session.execute(update(Key).where(Key.tg_id == tg_id, Key.client_id == client_id).values(**values)) logger.info(f"[Store Key] Ключ обновлён: tg_id={tg_id}, client_id={client_id}, server_id={server_id}") else: if current_device_limit is None: current_device_limit = selected_device_limit if current_traffic_limit is None: current_traffic_limit = selected_traffic_limit new_key = Key( tg_id=tg_id, client_id=client_id, email=email, created_at=int(datetime.utcnow().timestamp() * 1000), expiry_time=expiry_time, key=key, server_id=server_id, remnawave_link=remnawave_link, tariff_id=tariff_id, alias=alias, selected_device_limit=selected_device_limit, selected_traffic_limit=selected_traffic_limit, selected_price_rub=selected_price_rub, current_device_limit=current_device_limit, current_traffic_limit=current_traffic_limit, ) session.add(new_key) logger.info(f"[Store Key] Ключ создан: tg_id={tg_id}, client_id={client_id}, server_id={server_id}") await session.commit() invalidate_user_snapshot(tg_id) except SQLAlchemyError as e: logger.error(f"❌ Ошибка при сохранении ключа: {e}") await session.rollback() raise async def get_keys(session: AsyncSession, tg_id: int): result = await session.execute(select(Key).where(Key.tg_id == tg_id)) return result.scalars().all() async def get_all_keys(session: AsyncSession): result = await session.execute(select(Key)) return result.scalars().all() async def get_key_by_server(session: AsyncSession, tg_id: int, client_id: str): stmt = select(Key).where(Key.tg_id == tg_id, Key.client_id == client_id) result = await session.execute(stmt) return result.scalar_one_or_none() async def get_key_details(session: AsyncSession, email: str) -> dict | None: """Возвращает подробную информацию о ключе по email.""" stmt = select(Key, User).join(User, Key.tg_id == User.tg_id).where(Key.email == email) result = await session.execute(stmt) row = result.first() if not row: return None key, user = row expiry_date = datetime.utcfromtimestamp(key.expiry_time / 1000) current_date = datetime.utcnow() time_left = expiry_date - current_date if time_left.total_seconds() <= 0: days_left_message = "Ключ истек." elif time_left.days > 0: days_left_message = f"Осталось дней: {time_left.days}" else: hours_left = time_left.seconds // 3600 days_left_message = f"Осталось часов: {hours_left}" return { "key": key.key, "remnawave_link": key.remnawave_link, "server_id": key.server_id, "created_at": key.created_at, "expiry_time": key.expiry_time, "client_id": key.client_id, "tg_id": user.tg_id, "email": key.email, "is_frozen": key.is_frozen, "balance": user.balance, "alias": key.alias, "expiry_date": expiry_date.strftime("%d %B %Y года %H:%M"), "days_left_message": days_left_message, "link": key.key or key.remnawave_link, "cluster_name": key.server_id, "location_name": key.server_id, "tariff_id": key.tariff_id, "selected_device_limit": key.selected_device_limit, "selected_traffic_limit": key.selected_traffic_limit, "selected_price_rub": key.selected_price_rub, "current_device_limit": key.current_device_limit, "current_traffic_limit": key.current_traffic_limit, } async def get_key_count(session: AsyncSession, tg_id: int) -> int: result = await session.execute(select(func.count()).select_from(Key).where(Key.tg_id == tg_id)) return result.scalar() or 0 async def delete_key(session: AsyncSession, identifier: int | str, commit: bool = True): tg_id_for_cache = None if isinstance(identifier, str): res = await session.execute(select(Key.tg_id).where(Key.client_id == identifier).limit(1)) tg_id_for_cache = res.scalar_one_or_none() else: tg_id_for_cache = identifier stmt = delete(Key).where(Key.tg_id == identifier if isinstance(identifier, int) else Key.client_id == identifier) await session.execute(stmt) if commit: await session.commit() if tg_id_for_cache is not None: invalidate_user_snapshot(tg_id_for_cache) logger.info(f"Ключ с идентификатором {identifier} удалён") async def update_key_expiry(session: AsyncSession, client_id: str, new_expiry_time: int): await session.execute(update(Key).where(Key.client_id == client_id).values(expiry_time=new_expiry_time)) await session.commit() logger.info(f"Срок действия ключа {client_id} обновлён до {new_expiry_time}") async def get_client_id_by_email(session: AsyncSession, email: str): result = await session.execute(select(Key.client_id).where(Key.email == email)) return result.scalar_one_or_none() async def update_key_notified(session: AsyncSession, tg_id: int, client_id: str): await session.execute(update(Key).where(Key.tg_id == tg_id, Key.client_id == client_id).values(notified=True)) await session.commit() async def mark_key_as_frozen(session: AsyncSession, tg_id: int, client_id: str, time_left: int): await session.execute( text( """ UPDATE keys SET expiry_time = :expiry, is_frozen = TRUE WHERE tg_id = :tg_id AND client_id = :client_id """ ), {"expiry": time_left, "tg_id": tg_id, "client_id": client_id}, ) async def mark_key_as_unfrozen( session: AsyncSession, tg_id: int, client_id: str, new_expiry_time: int, ): await session.execute( text( """ UPDATE keys SET expiry_time = :expiry, is_frozen = FALSE WHERE tg_id = :tg_id AND client_id = :client_id """ ), {"expiry": new_expiry_time, "tg_id": tg_id, "client_id": client_id}, ) async def update_key_tariff(session: AsyncSession, client_id: str, tariff_id: int): await session.execute(update(Key).where(Key.client_id == client_id).values(tariff_id=tariff_id)) await session.commit() logger.info(f"Тариф ключа {client_id} обновлён на {tariff_id}") async def get_subscription_link(session: AsyncSession, email: str) -> str | None: result = await session.execute(select(func.coalesce(Key.key, Key.remnawave_link)).where(Key.email == email)) return result.scalar_one_or_none() async def update_key_client_id(session: AsyncSession, email: str, new_client_id: str): await session.execute(update(Key).where(Key.email == email).values(client_id=new_client_id)) await session.commit() logger.info(f"client_id обновлён для {email} -> {new_client_id}") async def update_key_link(session: AsyncSession, email: str, link: str) -> bool: q = update(Key).where(Key.email == email).values(key=link).returning(Key.client_id) res = await session.execute(q) await session.commit() return res.scalar_one_or_none() is not None async def save_key_config_with_mode( session: AsyncSession, email: str, selected_devices: int | None, selected_traffic_gb: int | None, total_price: int, has_device_choice: bool, has_traffic_choice: bool, config_mode: str, ) -> None: values: dict = {} if config_mode == "pack": if has_device_choice and selected_devices is not None: values["current_device_limit"] = int(selected_devices) if has_traffic_choice and selected_traffic_gb is not None: values["current_traffic_limit"] = int(selected_traffic_gb) else: device_val = int(selected_devices) if selected_devices is not None and has_device_choice else None traffic_val = int(selected_traffic_gb) if selected_traffic_gb is not None and has_traffic_choice else None values["selected_device_limit"] = device_val values["selected_traffic_limit"] = traffic_val values["selected_price_rub"] = int(total_price) values["current_device_limit"] = device_val values["current_traffic_limit"] = traffic_val if not values: return await session.execute(update(Key).where(Key.email == email).values(**values)) async def reset_key_current_limits_to_selected(session: AsyncSession, client_id: str): """Сбрасывает текущие лимиты к выбранным для ключа.""" await session.execute( text( """ UPDATE keys SET current_device_limit = selected_device_limit, current_traffic_limit = selected_traffic_limit WHERE client_id = :client_id """ ), {"client_id": client_id}, ) await session.commit() logger.info(f"Текущие лимиты ключа {client_id} сброшены к выбранным")