umb/bot/routers/ai.py

import logging
import time

from aiogram import Router, F
from aiogram.types import Message
from aiogram.filters import Command

from bot.utils.ai_client import ask_ai
from bot.utils.database import (
    add_context_message,
    get_user_context,
    is_ai_blocked,
    block_user_from_ai,
    unblock_user_from_ai,
    get_chat_users,
    clear_dialogue,
    clear_user_context,
    get_or_create_dialogue,
    save_chat_user,
)
from bot.utils.memory import (
    save_summary_with_embedding,
    find_relevant_summaries,
    generate_summary,
)
from config import (
    AI_BLOCK_DEFAULT_DURATION,
    AI_CONTEXT_LIMIT,
    AI_DIALOGUE_LIMIT,
    AI_PHASE2_LIMIT,
)

logger = logging.getLogger(__name__)

router = Router()


def extract_username(text: str) -> str | None:
    parts = text.split()
    for part in parts:
        if part.startswith("@"):
            return part[1:]
        if part.startswith("id") and part[2:].isdigit():
            return part[2:]
    return None


async def get_user_id_by_username(message: Message, username: str) -> int | None:
    try:
        username_clean = username.lstrip("@")
        admins = await message.bot.get_chat_administrators(message.chat.id)
        for admin in admins:
            if admin.user.username and admin.user.username.lower() == username_clean.lower():
                return admin.user.id
            if str(admin.user.id) == username_clean:
                return admin.user.id

        if username_clean.isdigit():
            return int(username_clean)
    except Exception as exc:
        logger.error("Error resolving username %s: %s", username, exc)

    return None


async def _is_creator(message: Message) -> bool:
    if message.from_user.id == message.chat.id:
        return True
    try:
        admins = await message.bot.get_chat_administrators(message.chat.id)
        for admin in admins:
            if admin.status == "creator" and admin.user.id == message.from_user.id:
                return True
    except Exception as exc:
        logger.error("Error checking creator: %s", exc)
    return False


async def _not_creator(message: Message) -> bool:
    return not await _is_creator(message)


@router.message(Command("ai"))
async def cmd_ai(message: Message):
    user_id = message.from_user.id
    chat_id = message.chat.id

    await save_chat_user(
        user_id,
        chat_id,
        message.from_user.username,
        message.from_user.full_name or "",
    )

    if await is_ai_blocked(user_id, chat_id):
        await message.answer(
            "Тебе временно недоступен AI. Обратись к владельцу чата.",
            parse_mode=None,
        )
        return

    text_parts = message.text.split(maxsplit=1)
    has_direct_question = len(text_parts) > 1 and text_parts[1].strip()

    if message.reply_to_message:
        target_text = message.reply_to_message.text or message.reply_to_message.caption
        if not target_text:
            await message.answer(
                "Могу работать только с текстовыми сообщениями.",
                parse_mode=None,
            )
            return
        prompt = f"Проанализируй это сообщение:\n\n{target_text}"
        context_text = target_text
    elif message.forward_origin:
        target_text = message.text or message.caption
        if not target_text:
            await message.answer(
                "Могу работать только с текстовыми сообщениями.",
                parse_mode=None,
            )
            return
        prompt = f"Проанализируй это сообщение:\n\n{target_text}"
        context_text = target_text
    elif has_direct_question:
        prompt = text_parts[1].strip()
        context_text = prompt
    else:
        await _start_dialogue(message)
        return

    status_msg = await message.reply("Думаю...", parse_mode=None)

    context = await get_user_context(user_id, chat_id, AI_CONTEXT_LIMIT)

    extra_system_content = None
    try:
        relevant = await find_relevant_summaries(user_id, chat_id, prompt, top_k=2)
        if relevant:
            extra_system_content = "Из прошлых диалогов:\n" + "\n---\n".join(relevant)
    except Exception as exc:
        logger.error("Error fetching relevant summaries: %s", exc)

    async def update_status(text: str):
        try:
            await status_msg.edit_text(text, parse_mode=None)
        except Exception as exc:
            logger.debug("Status update failed: %s", exc)

    response = await ask_ai(
        prompt,
        context,
        status_callback=update_status,
        extra_system_content=extra_system_content,
    )

    if not response:
        response = "Не удалось получить ответ от AI."

    await add_context_message(user_id, chat_id, context_text, "user")
    await add_context_message(user_id, chat_id, response, "assistant")

    try:
        await status_msg.edit_text(
            response, parse_mode="HTML", disable_web_page_preview=True
        )
    except Exception:
        try:
            await status_msg.edit_text(response, parse_mode=None, disable_web_page_preview=True)
        except Exception as exc:
            logger.warning("Failed to edit status message: %s", exc)


async def _start_dialogue(message: Message):
    user_id = message.from_user.id
    chat_id = message.chat.id

    dialogue = await get_or_create_dialogue(user_id, chat_id)

    if not dialogue["is_active"]:
        remaining = (
            int(dialogue["blocked_until"] - time.time()) if dialogue["blocked_until"] else 0
        )
        mins = max(0, remaining // 60)
        await message.answer(
            f"Твой лимит диалога исчерпан. Попробуй через {mins} мин.",
            parse_mode=None,
        )
        return

    if dialogue["phase"] == 2 and dialogue["msg_count"] >= AI_PHASE2_LIMIT:
        await message.answer("Твой лимит исчерпан.", parse_mode=None)
        return

    await message.answer(
        "💬 Диалог начат! Отвечай на мои сообщения, чтобы продолжать.\n"
        "Отправь /aiclear чтобы завершить.",
        parse_mode=None,
    )


@router.message(Command("aiuser"))
async def cmd_aiuser(message: Message):
    if await _not_creator(message):
        return

    try:
        users = await get_chat_users(message.chat.id)

        lines = []
        for i, u in enumerate(users, 1):
            username_str = f"@{u['username']}" if u.get("username") else "NO DATA"
            lines.append(f"{i}. {u['full_name']} | {username_str} | {u['user_id']}")

        text = "\n".join(lines) if lines else "Нет данных о пользователях."

        await message.bot.send_message(
            message.from_user.id,
            f"📋 Список пользователей чата ({len(users)}):\n\n{text}",
            parse_mode=None,
        )

        await message.answer("✅ Список отправлен в ЛС.", parse_mode=None)
    except Exception as exc:
        logger.error("aiuser error: %s", exc)
        await message.answer(
            "Не удалось получить список пользователей. Возможно, у бота нет доступа.",
            parse_mode=None,
        )


@router.message(Command("aiclear"))
async def cmd_aiclear(message: Message):
    user_id = message.from_user.id
    chat_id = message.chat.id

    try:
        context_messages = await get_user_context(user_id, chat_id, limit=50)
        cnt = len(context_messages)
        if cnt >= 2:
            status_msg = await message.answer(
                "Сохраняю выжимку диалога...", parse_mode=None
            )
            summary = await generate_summary(context_messages)
            if summary:
                await save_summary_with_embedding(user_id, chat_id, summary)
                await clear_user_context(user_id, chat_id)
                await status_msg.edit_text(
                    "✅ Выжимка сохранена. Диалог очищен.", parse_mode=None
                )
                logger.info(
                    "aiclear | user=%d chat=%d context=%d summary=%d",
                    user_id,
                    chat_id,
                    cnt,
                    len(summary),
                )
            else:
                logger.warning(
                    "aiclear | user=%d chat=%d context=%d summary=None",
                    user_id,
                    chat_id,
                    cnt,
                )
        else:
            logger.warning(
                "aiclear | user=%d chat=%d not enough context (%d < 2)",
                user_id,
                chat_id,
                cnt,
            )
    except Exception as exc:
        logger.error("aiclear summary error: %s", exc)

    await clear_dialogue(user_id, chat_id)
    await unblock_user_from_ai(user_id, chat_id)

    await message.answer(
        "✅ Диалог очищен. Можешь начать новый.",
        parse_mode=None,
    )


@router.message(Command("aino"))
async def cmd_aino(message: Message):
    if await _not_creator(message):
        return

    username = extract_username(message.text)
    if not username:
        await message.answer(
            "Укажи пользователя: /aino @username или /aino id<UID>",
            parse_mode=None,
        )
        return

    target_user_id = await get_user_id_by_username(message, username)
    if not target_user_id:
        await message.answer(
            f"Не удалось найти пользователя {username}.",
            parse_mode=None,
        )
        return

    await block_user_from_ai(
        target_user_id,
        message.chat.id,
        message.from_user.id,
        AI_BLOCK_DEFAULT_DURATION,
    )
    await message.answer(
        f"Пользователь {username} заблокирован от AI на 24 часа.",
        parse_mode=None,
    )


@router.message(Command("aiyes"))
async def cmd_aiyes(message: Message):
    if await _not_creator(message):
        return

    username = extract_username(message.text)
    if not username:
        await message.answer(
            "Укажи пользователя: /aiyes @username или /aiyes id<UID>",
            parse_mode=None,
        )
        return

    target_user_id = await get_user_id_by_username(message, username)
    if not target_user_id:
        await message.answer(
            f"Не удалось найти пользователя {username}.",
            parse_mode=None,
        )
        return

    unblocked = await unblock_user_from_ai(target_user_id, message.chat.id)
    if unblocked:
        await message.answer(
            f"Пользователь {username} разблокирован для AI.",
            parse_mode=None,
        )
    else:
        await message.answer(
            f"Пользователь {username} не был заблокирован от AI.",
            parse_mode=None,
        )