from functools import partial

from aiogram import Router, F, Bot
from aiogram.enums import ChatType
from aiogram.filters import Command, CommandObject, CommandStart
from aiogram.types import Message

import ai
import utils
from messages import *

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

router = Router()


@router.message(CommandStart(), F.chat.type == ChatType.PRIVATE)
async def start_handler(message: Message, bot: Bot):
    chat_id = message.chat.id
    database.DB.create_chat_if_not_exists(bot.id, chat_id)
    database.DB.chat_update(bot.id, chat_id, active=1)
    await message.answer("Привет!")


@router.message(Command("личность", prefix="!"), F.chat.type == ChatType.PRIVATE)
async def set_prompt_handler(message: Message, command: CommandObject, bot: Bot):
    chat_id = message.chat.id
    database.DB.create_chat_if_not_exists(bot.id, chat_id)

    database.DB.chat_update(bot.id, chat_id, ai_prompt=command.args)
    await message.answer("Личность ИИ изменена.")


@router.message(Command("сброс", prefix="!"), F.chat.type == ChatType.PRIVATE)
async def reset_context_handler(message: Message, bot: Bot):
    chat_id = message.chat.id
    database.DB.create_chat_if_not_exists(bot.id, chat_id)

    ai.agent_instance.clear_chat_context(bot.id, chat_id)
    await message.answer("Контекст очищен.")


@router.message(F.content_type.in_(ACCEPTED_CONTENT_TYPES), F.chat.type == ChatType.PRIVATE)
async def any_message_handler(message: Message, bot: Bot):
    chat_id = message.chat.id

    try:
        ai_message = await create_ai_message(message, bot)
    except utils.UnsupportedContentType:
        await message.answer(MESSAGE_UNSUPPORTED_CONTENT_TYPE)
        return

    ai_fwd_messages: list[ai.Message] = []
    if message.reply_to_message:
        last_id = ai.agent_instance.get_last_assistant_message_id(bot.id, chat_id)
        if message.reply_to_message.message_id != last_id:
            # Оригинального сообщения нет в контексте, или оно не последнее -> переслать его
            fwd_message = await create_ai_message(message.reply_to_message, bot)
            if (message.reply_to_message.from_user is not None and message.from_user is not None and
                    message.reply_to_message.from_user.id == message.from_user.id):
                # Замаскировать реальное имя пользователя
                fwd_message.user_name = "Пользователь"
            ai_fwd_messages = [fwd_message]

    answer: ai.Message
    success: bool
    answer, success = await utils.run_with_progress(
        partial(ai.agent_instance.get_private_chat_reply, bot.id, chat_id, ai_message, ai_fwd_messages),
        partial(bot.send_chat_action, chat_id, 'typing'),
        interval=4)

    if answer.image is not None and answer.image_hires is not None:
        answer_id = (await message.answer_photo(photo=wrap_photo(answer.image),
                                                caption=trim_caption(answer.text))).message_id
        await message.answer_document(document=wrap_document(answer.image_hires, 'image', 'png'))
    elif answer.text is not None:
        answer_id = (await message.answer(answer.text)).message_id
    else:
        return

    if success:
        ai.agent_instance.set_last_response_id(bot.id, chat_id, answer_id)