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 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() 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): stmt = delete(Key).where(Key.tg_id == identifier if str(identifier).isdigit() else Key.client_id == identifier) await session.execute(stmt) if commit: await session.commit() 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} сброшены к выбранным")