import asyncio

from aiogram import Router, Bot, F
from aiogram.enums import ParseMode
from aiogram.filters import *
from aiogram.types import *

from adapters.google_adapter import GoogleAdapter

router = Router()

@router.message(Command("image"))
async def cmd_image_gen(message: Message, command: CommandObject, gemini: GoogleAdapter, bot: Bot):
    # ... ваш код ...
    # Обратите внимание: gemini теперь прилетает аргументом, так как мы сделали dp["gemini"]
    prompt = command.args
    if not prompt:
        await message.answer("⚠️ Напиши промпт.")
        return

    wait_msg = await message.answer("🎨 Генерирую...")

    # Получение байтов фото (логика та же)
    image_bytes = None
    if message.photo:
        file_io = await bot.download(message.photo[-1].file_id)
        image_bytes = file_io.getvalue()
    elif message.reply_to_message and message.reply_to_message.photo:
        file_io = await bot.download(message.reply_to_message.photo[-1].file_id)
        image_bytes = file_io.getvalue()

    result = await asyncio.to_thread(
        gemini.generate, prompt=prompt, image_bytes=image_bytes, generate_image=True
    )

    await wait_msg.delete()

    if result.get("images"):
        for img in result["images"]:
            await message.answer_document(BufferedInputFile(img.read(), "img.png"))
    elif result.get("text"):
        await message.answer(result["text"])
    else:
        await message.answer(f"Ошибка: {result.get('error', 'Unknown')}")


@router.message(F.text)
async def handle_text(message: Message, gemini: GoogleAdapter, bot: Bot):
    await bot.send_chat_action(message.chat.id, "typing")
    result = await asyncio.to_thread(gemini.generate, prompt=message.text)
    if result.get("text"):
        await message.answer(result["text"], parse_mode=ParseMode.MARKDOWN)
    else:
        await message.answer("Ошибка генерации")