from functools import partial

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

import ai_agent
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.agent.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

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

    if answer.image is not None:
        answer_id = (await message.answer_photo(photo=wrap_photo(answer.image), caption=answer.text)).message_id
        await message.answer_document(document=wrap_document(answer.image_hires, 'image', 'png'))
    else:
        answer_id = (await message.answer(answer.text)).message_id
    if success:
        ai_agent.agent.set_last_response_id(bot.id, chat_id, answer_id)