from types import UnionType from aiogram import Router, F from aiogram.fsm.context import FSMContext from aiogram.types import Message, CallbackQuery, InlineKeyboardButton from aiogram.utils.keyboard import InlineKeyboardBuilder from logging import getLogger from sqlmodel.ext.asyncio.session import AsyncSession from typing import get_args, get_origin, TYPE_CHECKING from ....model.bot_entity import BotEntity from ....model.owned_bot_entity import OwnedBotEntity from ....model.bot_enum import BotEnum from ....model.settings import Settings from ....model.user import UserBase from ....model.view_setting import ViewSetting from ....model.descriptors import EntityFieldDescriptor, EntityDescriptor from ....model import EntityPermission from ....utils import serialize, deserialize, get_user_permissions from ..context import ContextData, CallbackCommand from ..common import (get_send_message, get_local_text, get_field_descriptor, get_entity_descriptor, add_pagination_controls, add_filter_controls) from .common import wrap_editor if TYPE_CHECKING: from ....main import QBotApp logger = getLogger(__name__) router = Router() async def entity_picker(message: Message | CallbackQuery, field_descriptor: EntityFieldDescriptor, edit_prompt: str, current_value: BotEntity | BotEnum | list[BotEntity] | list[BotEnum], **kwargs): state_data: dict = kwargs["state_data"] state_data.update({"current_value": serialize(current_value, field_descriptor), "value": serialize(current_value, field_descriptor), "edit_prompt": edit_prompt}) await render_entity_picker(field_descriptor = field_descriptor, message = message, current_value = current_value, edit_prompt = edit_prompt, **kwargs) def calc_total_pages(items_count: int, page_size: int) -> int: return max(items_count // page_size + (1 if items_count % page_size else 0), 1) async def render_entity_picker(*, field_descriptor: EntityFieldDescriptor, message: Message | CallbackQuery, callback_data: ContextData, user: UserBase, db_session: AsyncSession, state: FSMContext, current_value: BotEntity | BotEnum | list[BotEntity] | list[BotEnum], edit_prompt: str, page: int = 1, **kwargs): if callback_data.command in [CallbackCommand.ENTITY_PICKER_PAGE, CallbackCommand.ENTITY_PICKER_TOGGLE_ITEM]: page = int(callback_data.data.split("&")[0]) # is_list = False # type_origin = get_origin(field_descriptor.type_) # if type_origin == UnionType: # type_ = get_args(field_descriptor.type_)[0] # elif type_origin == list: # type_ = get_args(field_descriptor.type_)[0] # is_list = True # else: # type_ = field_descriptor.type_ type_ = field_descriptor.type_base is_list = field_descriptor.is_list if not issubclass(type_, BotEntity) and not issubclass(type_, BotEnum): raise ValueError("Unsupported type") page_size = await Settings.get(Settings.PAGE_SIZE) if issubclass(type_, BotEnum): items_count = len(type_.all_members) total_pages = calc_total_pages(items_count, page_size) page = min(page, total_pages) enum_items = list(type_.all_members.values())[page_size * (page - 1):page_size * page] items = [{"text": f"{"" if not is_list else "【✔︎】 " if item in (current_value or []) else "【 】 "}{item.localized(user.lang)}", "value": item.value} for item in enum_items] else: permissions = get_user_permissions(user, type_.bot_entity_descriptor) entity_filter = await ViewSetting.get_filter(session = db_session, user_id = user.id, entity_name = type_.bot_entity_descriptor.class_name) if (EntityPermission.LIST_ALL in permissions or (EntityPermission.LIST in permissions and not issubclass(type_, OwnedBotEntity))): items_count = await type_.get_count(session = db_session, filter = entity_filter) total_pages = calc_total_pages(items_count, page_size) page = min(page, total_pages) entity_items = await type_.get_multi( session = db_session, order_by = type_.name, filter = entity_filter, skip = page_size * (page - 1), limit = page_size) elif (EntityPermission.LIST in permissions and issubclass(type_, OwnedBotEntity)): items_count = await type_.get_count_by_user(session = db_session, user_id = user.id, filter = entity_filter) total_pages = calc_total_pages(items_count, page_size) page = min(page, total_pages) entity_items = await type_.get_multi_by_user( session = db_session, user_id = user.id, order_by = type_.name, filter = entity_filter, skip = page_size * (page - 1), limit = page_size) else: items_count = 0 total_pages = 1 page = 1 entity_items = list[BotEntity]() items = [{"text": f"{"" if not is_list else "【✔︎】 " if item in (current_value or []) else "【 】 "}{ type_.bot_entity_descriptor.item_caption(type_.bot_entity_descriptor, item) if type_.bot_entity_descriptor.item_caption else get_local_text(item.name, user.lang) if type_.bot_entity_descriptor.fields_descriptors["name"].localizable else item.name}", "value": str(item.id)} for item in entity_items] # total_pages = items_count // page_size + (1 if items_count % page_size else 0) keyboard_builder = InlineKeyboardBuilder() for item in items: keyboard_builder.row( InlineKeyboardButton(text = item["text"], callback_data = ContextData( command = CallbackCommand.ENTITY_PICKER_TOGGLE_ITEM if is_list else CallbackCommand.FIELD_EDITOR_CALLBACK, context = callback_data.context, entity_name = callback_data.entity_name, entity_id = callback_data.entity_id, field_name = callback_data.field_name, data = f"{page}&{item['value']}" if is_list else item["value"], save_state = True).pack())) add_pagination_controls(keyboard_builder = keyboard_builder, callback_data = callback_data, total_pages = total_pages, command = CallbackCommand.ENTITY_PICKER_PAGE, page = page) if issubclass(type_, BotEntity): add_filter_controls(keyboard_builder = keyboard_builder, entity_descriptor = type_.bot_entity_descriptor, filter = entity_filter) if is_list: keyboard_builder.row( InlineKeyboardButton(text = await Settings.get(Settings.APP_STRINGS_DONE_BTN), callback_data = ContextData( command = CallbackCommand.FIELD_EDITOR_CALLBACK, context = callback_data.context, entity_name = callback_data.entity_name, entity_id = callback_data.entity_id, field_name = callback_data.field_name, save_state = True).pack())) state_data = kwargs["state_data"] await wrap_editor(keyboard_builder = keyboard_builder, field_descriptor = field_descriptor, callback_data = callback_data, state_data = state_data) await state.set_data(state_data) send_message = get_send_message(message) await send_message(text = edit_prompt, reply_markup = keyboard_builder.as_markup()) @router.callback_query(ContextData.filter(F.command == CallbackCommand.ENTITY_PICKER_PAGE)) @router.callback_query(ContextData.filter(F.command == CallbackCommand.ENTITY_PICKER_TOGGLE_ITEM)) async def entity_picker_callback(query: CallbackQuery, callback_data: ContextData, db_session: AsyncSession, app: "QBotApp", state: FSMContext, **kwargs): state_data = await state.get_data() kwargs["state_data"] = state_data field_descriptor = get_field_descriptor(app = app, callback_data = callback_data) current_value = await deserialize(session = db_session, type_ = field_descriptor.type_, value = state_data["current_value"]) edit_prompt = state_data["edit_prompt"] value = await deserialize(session = db_session, type_ = field_descriptor.type_, value = state_data["value"]) if callback_data.command == CallbackCommand.ENTITY_PICKER_TOGGLE_ITEM: page, id_value = callback_data.data.split("&") page = int(page) type_ = field_descriptor.type_base if issubclass(type_, BotEnum): item = type_(id_value) if item in value: value.remove(item) else: value.append(item) else: item = await type_.get(session = db_session, id = int(id_value)) if item in value: value.remove(item) else: value.append(item) state_data.update({"value": serialize(value, field_descriptor)}) elif callback_data.command == CallbackCommand.ENTITY_PICKER_PAGE: if callback_data.data == "skip": return page = int(callback_data.data) else: raise ValueError("Unsupported command") await render_entity_picker(field_descriptor = field_descriptor, message = query, callback_data = callback_data, current_value = value, edit_prompt = edit_prompt, db_session = db_session, app = app, state = state, page = page, **kwargs)