Add new_invoice command

3 files changed, 116 insertions, 4 deletions

diff --git a/handlers/admin/__init__.py b/handlers/admin/__init__.py
index 2f7c74f..98b127f 100644
--- a/handlers/admin/__init__.py
+++ b/handlers/admin/__init__.py
@@ -3,6 +3,7 @@ from aiogram.filters import MagicData
 
 # isort: off
 from . import new_announcement
+from . import new_invoice
 # isort: on
 
 router = Router(name="admin")
@@ -11,4 +12,5 @@ router.callback_query.filter(MagicData(F.user.is_admin()))
 
 router.include_routers(
     new_announcement.router,
+    new_invoice.router,
 )
diff --git a/handlers/admin/new_announcement.py b/handlers/admin/new_announcement.py
index 79cf8d4..0920c47 100644
--- a/handlers/admin/new_announcement.py
+++ b/handlers/admin/new_announcement.py
@@ -1,6 +1,7 @@
 from datetime import UTC, datetime
 
 from aiogram import Bot, F, Router
+from aiogram.enums import ButtonStyle
 from aiogram.exceptions import TelegramAPIError
 from aiogram.filters import Command
 from aiogram.fsm.context import FSMContext
@@ -17,7 +18,7 @@ from sqlalchemy import select
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from libs.fsm import get_data, set_data
-from libs.msg import publish
+from libs.msg import publish_announcement
 from models import Announcement, RichText, User
 
 router = Router(name="new_announcement")
@@ -42,8 +43,8 @@ async def new_announcement_command(msg: Message, state: FSMContext) -> None:
         reply_markup=ReplyKeyboardMarkup(
             keyboard=[
                 [
-                    KeyboardButton(text=SEND_BUTTON),
-                    KeyboardButton(text=CANCEL_BUTTON),
+                    KeyboardButton(text=SEND_BUTTON, style=ButtonStyle.SUCCESS),
+                    KeyboardButton(text=CANCEL_BUTTON, style=ButtonStyle.DANGER),
                 ]
             ],
             resize_keyboard=True,
@@ -69,7 +70,7 @@ async def announcement_send(
     status_template = "Публикация анонса...\nОпубликовано: {}"
     status_msg = await msg.answer(status_template.format(0))
 
-    async for n in publish(bot, users, data.rich_text):
+    async for n in publish_announcement(bot, users, data.rich_text):
         try:
             await status_msg.edit_text(status_template.format(n))
         except TelegramAPIError:
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)