from functools import partial

from aiogram import Router, F, Bot
from aiogram.types import Message
from aiogram.enums.content_type import ContentType

import ai_agent
import utils
from messages import *

import tg.tg_database as database
from tg.utils import *

router = Router()

ACCEPTED_CONTENT_TYPES: list[ContentType] = [
    ContentType.TEXT,
    ContentType.ANIMATION,
    ContentType.AUDIO,
    ContentType.DOCUMENT,
    ContentType.PAID_MEDIA,
    ContentType.PHOTO,
    ContentType.STICKER,
    ContentType.STORY,
    ContentType.VIDEO,
    ContentType.VIDEO_NOTE,
    ContentType.VOICE,
    ContentType.CHECKLIST,
    ContentType.CONTACT,
    ContentType.POLL,
    ContentType.VENUE,
    ContentType.LOCATION
]


@router.message(F.content_type.in_(ACCEPTED_CONTENT_TYPES))
async def any_message_handler(message: Message, bot: Bot):
    chat_id = message.chat.id
    chat = database.DB.create_chat_if_not_exists(bot.id, chat_id)
    if chat['active'] == 0:
        return

    # Игнорировать ботов
    if message.from_user.is_bot:
        return

    user_id = message.from_user.id
    database.DB.create_user_if_not_exists(bot.id, chat_id, user_id)
    database.DB.user_set_last_message(bot.id, chat_id, user_id, utils.posix_time())
    database.DB.user_increment_messages(bot.id, chat_id, user_id)

    bot_user = await bot.me()

    ai_fwd_messages: list[ai_agent.Message] = []

    try:
        message_text = get_message_text(message)
        bot_username_mention = '@' + bot_user.username
        if message_text is not None and message_text.find(bot_username_mention) != -1:
            # Сообщение содержит @bot_username
            message_text = message_text.replace(bot_username_mention, bot_user.first_name)
            if message.reply_to_message:
                # Сообщение также является ответом -> переслать оригинальное сообщение
                ai_fwd_messages = [await create_ai_message(message.reply_to_message, bot)]
        elif message.reply_to_message and message.reply_to_message.from_user.id == bot_user.id:
            # Ответ на сообщение бота
            last_id = ai_agent.agent.get_last_assistant_message_id(bot.id, chat_id)
            if message.reply_to_message.message_id != last_id:
                # Оригинального сообщения нет в контексте, или оно не последнее -> переслать его
                ai_fwd_messages = [await create_ai_message(message.reply_to_message, bot)]
        else:
            return
    except UnsupportedContentType:
        await message.reply(MESSAGE_UNSUPPORTED_CONTENT_TYPE)
        return

    ai_message = await create_ai_message(message, bot)
    ai_message.text = message_text

    answer, success = await utils.run_with_progress(
            partial(ai_agent.agent.get_group_chat_reply, bot.id, chat_id, ai_message, ai_fwd_messages),
            partial(message.bot.send_chat_action, chat_id, 'typing'),
            interval=4)

    answer_id = (await message.reply(answer)).message_id
    if success:
        ai_agent.agent.set_last_response_id(bot.id, chat_id, answer_id)