from typing import Dict, List, Optional

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

from openrouter import OpenRouter
from dataclasses import dataclass

import utils
import tg.tg_database as database

router = Router()


@dataclass
class ChatContext:
    def __init__(self, max_messages: int):
        self.max_messages: int = max_messages
        self.messages: List[Dict[str, str]] = []

    def add_message(self, role: str, content: str):
        if len(self.messages) == self.max_messages:
            # Всегда сохраняем в контексте системное сообщение
            self.messages.pop(1)
        self.messages.append({"role": role, "content": content})

    def get_messages_for_api(self) -> List[Dict[str, str]]:
        return self.messages


chat_contexts: Dict[int, ChatContext] = {}
bot_user: Optional[User] = None

system_prompt = """
Ты - помощник в групповом чате Telegram.
Отвечай на вопросы и поддерживай контекст беседы.
Ты не можешь обсуждать политику и религию.
Сообщения пользователей будут приходить в следующем формате: '[Имя]: текст сообщения'
При ответе НЕ нужно указывать пользователя, которому он предназначен.
"""


def get_ai_chat_context(chat_id: int) -> ChatContext:
    """Get or create chat context for a specific chat"""
    if chat_id not in chat_contexts:
        chat_contexts[chat_id] = ChatContext(10)
        chat = database.DB.get_chat(chat_id)
        prompt = system_prompt
        if chat['ai_prompt'] is not None:
            prompt += '\n\n' + chat['ai_prompt']
        chat_contexts[chat_id].add_message(role="system", content=prompt)
    return chat_contexts[chat_id]


async def ai_message_handler(message: Message):
    chat_id = message.chat.id

    # Extract user information and message content
    if message.from_user.first_name and message.from_user.last_name:
        user_name = "{} {}".format(message.from_user.first_name, message.from_user.last_name)
    elif message.from_user.first_name:
        user_name = message.from_user.first_name
    elif message.from_user.username:
        user_name = message.from_user.username
    else:
        user_name = str(message.from_user.id)

    bot_mention = '@' + bot_user.username
    message_text = message.text.replace(bot_mention, bot_user.first_name)

    context = get_ai_chat_context(chat_id)
    context.add_message(
        role="user",
        content=f"[{user_name}]: {message_text}"
    )

    # noinspection PyUnresolvedReferences
    api_key = message.bot.config['openrouter_token']

    client = OpenRouter(api_key=api_key)
    messages_for_api = context.get_messages_for_api()

    await message.bot.send_chat_action(chat_id, 'typing')

    try:
        # Get response from OpenRouter
        response = await client.chat.send_async(
            model="meta-llama/llama-3.3-70b-instruct:free",
            messages=messages_for_api,
            max_tokens=500,
            temperature=0.7
        )

        # Extract AI response
        ai_response = response.choices[0].message.content

        # Add AI response to context
        context.add_message(role="assistant", content=ai_response)

        # Send response back to chat
        await message.reply(ai_response)

    except Exception as e:
        print(f"Error processing message: {e}")
        await message.reply("Извините, при обработке запроса произошла ошибка.")


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):
    chat_id = message.chat.id
    chat = database.DB.create_chat_if_not_exists(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(chat_id, user_id)
    database.DB.user_set_last_message(chat_id, user_id, utils.posix_time())
    database.DB.user_increment_messages(chat_id, user_id)

    global bot_user
    if bot_user is None:
        bot_user = await message.bot.get_me()

    bot_mention = '@' + bot_user.username
    if message.content_type == ContentType.TEXT and message.text.find(bot_mention) != -1:
        await ai_message_handler(message)