Add new_invoice command

1 files changed, 109 insertions, 0 deletions

diff --git a/handlers/admin/new_invoice.py b/handlers/admin/new_invoice.py
new file mode 100644
index 0000000..7e1a64d
--- /dev/null
+++ b/handlers/admin/new_invoice.py
@@ -0,0 +1,109 @@
+from datetime import UTC, datetime
+
+from aiogram import Bot, F, Router
+from aiogram.enums import ButtonStyle, ContentType
+from aiogram.exceptions import TelegramAPIError
+from aiogram.filters import Command
+from aiogram.fsm.context import FSMContext
+from aiogram.fsm.state import State, StatesGroup
+from aiogram.types import (
+    KeyboardButton,
+    Message,
+    ReplyKeyboardMarkup,
+    ReplyKeyboardRemove,
+)
+from pydantic import BaseModel
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from libs.fsm import get_data, set_data
+from libs.msg import send_invoice
+from models import Invoice, RichText, User
+
+router = Router(name="new_invoice")
+
+
+class NewInvoiceStates(StatesGroup):
+    message = State()
+
+
+class NewInvoiceData(BaseModel):
+    rich_text: RichText | None = None
+
+
+CREATE_BUTTON = "Создать"
+CANCEL_BUTTON = "Отменить создание"
+
+
+@router.message(Command("new_invoice"))
+async def new_invoice_command(msg: Message, state: FSMContext) -> None:
+    await msg.answer(
+        "Укажите сообщение для создания счёта",
+        reply_markup=ReplyKeyboardMarkup(
+            keyboard=[
+                [
+                    KeyboardButton(text=CREATE_BUTTON, style=ButtonStyle.SUCCESS),
+                    KeyboardButton(text=CANCEL_BUTTON, style=ButtonStyle.DANGER),
+                ]
+            ],
+            resize_keyboard=True,
+        ),
+    )
+    await state.set_state(NewInvoiceStates.message)
+
+
+@router.message(NewInvoiceStates.message, F.text == CREATE_BUTTON)
+async def invoice_send(
+    msg: Message,
+    bot: Bot,
+    state: FSMContext,
+    session: AsyncSession,
+) -> None:
+    users = await session.scalars(select(User.id).where(User.id != msg.chat.id))
+    data = await get_data(state, NewInvoiceData)
+
+    if data.rich_text is None:
+        await msg.answer("Для создания счёта укажите сообщение.")
+        return
+
+    status_template = "Рассылка счёта...\nОтправлено: {}"
+    status_msg = await msg.answer(status_template.format(0))
+
+    invoice = Invoice(message=data, datetime=datetime.now(UTC))
+    session.add(invoice)
+    await session.flush()
+
+    async for n in send_invoice(bot, users, data.rich_text, invoice.id):
+        try:
+            await status_msg.edit_text(status_template.format(n))
+        except TelegramAPIError:
+            pass
+
+    await status_msg.delete()
+    await msg.answer(
+        "Счёт отправлен всем пользователям",
+        # reply_markup=ReplyKeyboardRemove(),
+    )
+    # await state.clear()
+
+
+@router.message(NewInvoiceStates.message, F.text == CANCEL_BUTTON)
+async def invoice_cancel(msg: Message, state: FSMContext) -> None:
+    await msg.answer("Создание счёта отменено", reply_markup=ReplyKeyboardRemove())
+    await state.clear()
+
+
+@router.message(NewInvoiceStates.message)
+async def invoice_message(msg: Message, bot: Bot, state: FSMContext) -> None:
+    if msg.content_type != ContentType.TEXT or msg.text is None:
+        await msg.answer(
+            "Неверный тип сообщения.\n"
+            "Бот поддерживает отправку только текстовых счетов."
+        )
+        return
+
+    rich_text = RichText.from_message(msg)
+    await set_data(state, NewInvoiceData(rich_text=rich_text))
+
+    msg_rich_text = RichText.from_text("Сообщение вашего счёта:\n", rich_text)
+    await msg_rich_text.send(bot, msg.chat.id)