ai-char-bot/routers/char_router.py

import io
import logging
import traceback

from aiogram.filters import Command
from aiogram.fsm.context import FSMContext
from aiogram.fsm.state import State, StatesGroup
from aiogram.types import *
from aiogram import Router, F, Bot

from models.Asset import Asset, AssetType, AssetContentType
from models.Character import Character
from repos.dao import DAO

router = Router()


class States(StatesGroup):
    char_wait_name = State()
    char_wait_bio = State()


@router.message(F.document, Command("add_char"))
async def add_char(message: Message, state: FSMContext, dao: DAO):
    if not message.document.mime_type.startswith("image/"):
        await message.answer("❌ Пожалуйста, отправьте файл изображения (png/jpg).")
        return
    await state.set_data({"photo": message.document.file_id})
    await state.set_state(States.char_wait_name)
    await message.answer("Кайф, теперь напиши ее имя")


@router.message(States.char_wait_name)
async def new_char_name(message: Message, state: FSMContext, dao: DAO):
    data = await state.get_data()
    data["name"] = message.text
    await state.set_data(data)
    await state.set_state(States.char_wait_bio)
    await message.answer("А теперь напиши био. Хоть чуть чуть.")


# 3. Хендлер био (Финал)
@router.message(States.char_wait_bio, F.text)
async def new_char_bio(message: Message, state: FSMContext, dao: DAO, bot: Bot):
    # Получаем все накопленные данные
    data = await state.get_data()
    file_id = data["photo"]
    name = data["name"]
    bio = message.text

    wait_msg = await message.answer("💾 Сохраняю персонажа...")

    try:
        # 1. Скачиваем файл (один раз)
        # TODO: Для больших файлов лучше использовать streaming или сохранять во временный файл
        file_io = await bot.download(file_id)
        file_bytes = file_io.read()
        
        # 2. Создаем Character (сначала без ассета, чтобы получить ID)
        char = Character(
            id=None,
            name=name,
            character_image_tg_id=None,
            character_image_doc_tg_id=file_id,
            character_bio=bio,
            created_by=str(message.from_user.id)
        )
        
        # Сохраняем, чтобы получить ID
        await dao.chars.add_character(char)
        
        # 3. Создаем Asset (связанный с персонажем)
        avatar_asset_id = await dao.assets.create_asset(
            Asset(
                name="avatar.png", 
                type=AssetType.UPLOADED, 
                content_type=AssetContentType.IMAGE, 
                linked_char_id=str(char.id), 
                data=file_bytes,
                tg_doc_file_id=file_id
            )
        )
        
        # 4. Обновляем персонажа ссылками на ассет
        char.avatar_asset_id = avatar_asset_id
        char.avatar_image = f"/api/assets/{avatar_asset_id}" # Формируем ссылку вручную или используем метод, если появится
        
        # Отправляем подтверждение
        photo_msg = await message.answer_photo(
            photo=BufferedInputFile(file_bytes, filename="char.jpg"),
            caption=(
                "🎉 <b>Персонаж создан!</b>\n\n"
                f"👤 <b>Имя:</b> {char.name}\n"
                f"📝 <b>Био:</b> {char.character_bio}"
            )
        )
        
        # Сохраняем TG ID фото (которое отправили как фото, а не документ)
        char.character_image_tg_id = photo_msg.photo[-1].file_id

        # Финальное обновление персонажа
        await dao.chars.update_char(char.id, char)
        
        await wait_msg.delete()
        file_io.close()

        # Сбрасываем состояние
        await state.clear()

    except Exception as e:
        logger.error(f"Error creating character: {e}")
        traceback.print_exc()
        await wait_msg.edit_text(f"❌ Ошибка при сохранении: {e}")


@router.message(Command("chars"))
async def get_chars(message: Message, state: FSMContext, dao: DAO):
    wait_msg = await message.answer("Ищем персонажей")
    chars = await dao.chars.get_all_characters()
    keyboards = []
    if len(chars) > 0:
        for char in chars:
            keyboards.append(InlineKeyboardButton(text=char.name, callback_data=f'char_info_{char.id}'))
        keyboard = InlineKeyboardMarkup(inline_keyboard=[keyboards])
    else:
        keyboard = InlineKeyboardMarkup(
            inline_keyboard=[[InlineKeyboardButton(text="Персонажей нет", callback_data=f'no_chars')]])
    await message.answer("Сейчас есть такие персонажи:", reply_markup=keyboard)
    await wait_msg.delete()


@router.callback_query(F.data.startswith("char_info_"))
async def get_char_info(callback_query: CallbackQuery, state: FSMContext, dao: DAO, bot: Bot):
    await callback_query.message.delete()
    wait_msg = await callback_query.message.answer("Ищем инфу о персонаже")
    char = await dao.chars.get_character(callback_query.data.split("_")[-1])
    if char is None:
        await callback_query.message.answer("Информация о персонаже не найдена")
        await get_chars(callback_query.message, state, dao)
        await wait_msg.delete()
        return
    keyboard = InlineKeyboardMarkup(inline_keyboard=[
        [InlineKeyboardButton(text="Запросить фото в документе", callback_data=f'char_photo_file_{char.id}')]])

    photo_msg = await callback_query.message.answer_photo(
        photo=char.character_image_tg_id,
        caption=f"👤 <b>Имя:</b> {char.name}\n"
                f"📝 <b>Био:</b> {char.character_bio}",
        reply_markup=keyboard)
    await wait_msg.delete()


@router.callback_query(F.data.startswith("char_photo_file"))
async def get_char_info_photo_file(callback_query: CallbackQuery, state: FSMContext, dao: DAO):
    char = await dao.chars.get_character(callback_query.data.split("_")[-1])
    await callback_query.message.answer_document(char.character_image_doc_tg_id)


# 4. Хендлер-помощник (если отправили команду без файла)
@router.message(Command("add_char"))
async def add_char_help(message: Message):
    await message.answer(
        "ℹ️ **Как добавить персонажа:**\n"
        "Прикрепите фото (файлом/документом) и добавьте в подпись команду `/add_char`."
    )