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", 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", 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, )