from functools import partial

from vkbottle.bot import Message
from vkbottle.dispatch.rules.base import RegexRule
from vkbottle.framework.labeler import BotLabeler

import ai_agent
import utils
from messages import *

import vk.vk_database as database

labeler = BotLabeler()


@labeler.private_message(text="!старт")
async def start_handler(message: Message):
    chat_id = message.peer_id
    database.DB.create_chat_if_not_exists(chat_id)
    database.DB.chat_update(chat_id, active=1)
    await message.answer("Привет!")


@labeler.private_message(RegexRule(r"^!личность ((?:.|\n)+)"))
async def set_prompt_handler(message: Message, match):
    chat_id = message.peer_id
    database.DB.create_chat_if_not_exists(chat_id)

    database.DB.chat_update(chat_id, ai_prompt=match[0])
    await message.answer("Личность ИИ изменена.")


@labeler.private_message(text="!сброс")
async def reset_context_handler(message: Message):
    chat_id = message.peer_id
    database.DB.create_chat_if_not_exists(chat_id)

    ai_agent.agent.clear_chat_context(chat_id)
    await message.answer("Контекст очищен.")


@labeler.private_message()
async def any_message_handler(message: Message):
    chat_id = message.peer_id

    if len(message.text) == 0:
        await message.answer(MESSAGE_NOT_TEXT)
        return

    answer, success = await utils.run_with_progress(
        partial(ai_agent.agent.get_private_chat_reply, chat_id, message.text, message.message_id),
        partial(message.ctx_api.messages.set_activity, peer_id=chat_id, type='typing'),
        interval=4)

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