LawBot/bot/handlers/client/main.py

import aiogram.types as types
from aiogram import F, Router
from aiogram.filters import StateFilter
from aiogram.fsm.context import FSMContext

from create_bot import orm, tz
from keyboards.reply_keyboards import (
    CATEGORY_BY_TEXT,
    USER_TYPE_BY_TEXT,
    USER_TYPE_LABEL_BY_CODE,
    get_back_to_menu_kb,
    get_categories_kb,
    get_consultation_actions_ikb,
    get_consultations_ikb,
    get_delete_profile_kb,
    get_main_menu_kb,
    get_profile_kb,
    get_user_types_kb,
)
from states.client_states import AskQuestionStates, MainStates, ProfileStates
from utils.rag_api import ask_rag_answer
from utils.text_tools import format_llm_answer_html, split_plain_text_chunks, to_html

from datetime import datetime, timezone


client_router = Router()

DAILY_CONSULTATION_LIMIT = 5
CONSULTATION_MESSAGE_LIMIT = 5


def get_today_start_utc() -> datetime:
    now_local = datetime.now(tz)
    start_local = now_local.replace(hour=0, minute=0, second=0, microsecond=0)
    return start_local.astimezone(timezone.utc)


def build_start_text() -> str:
    return (
        "Здравствуйте. Я юридический ИИ-консультант по законам РФ.\n\n"
        "Я могу помочь:\n"
        "— разобраться в ситуации;\n"
        "— найти применимые нормы закона;\n"
        "— объяснить их простыми словами;\n"
        "— дать базовый план действий.\n\n"
        "Ответ носит информационный характер и не заменяет консультацию юриста."
    )


def build_help_text() -> str:
    return (
        "ℹ️ Как пользоваться ботом\n\n"
        "1. Нажмите «Задать вопрос».\n"
        "2. Выберите категорию.\n"
        "3. Опишите ситуацию.\n"
        "4. Укажите регион РФ.\n"
        "5. Получите ответ с найденными источниками.\n\n"
        "Пример хорошего вопроса:\n"
        "«Работодатель не выплатил зарплату за апрель. Работаю официально, Москва. Что делать?»\n\n"
        "Бот не заменяет юриста и не гарантирует результат спора."
    )


def build_profile_text(user) -> str:
    return (
        "👤 Профиль\n\n"
        f"Страна: {to_html(user.country or 'Россия')}\n"
        f"Регион: {to_html(user.region or 'не указан')}\n"
        f"Тип пользователя: {to_html(USER_TYPE_LABEL_BY_CODE.get(user.user_type, 'Физлицо'))}"
    )


def build_consultations_text(consultations: list) -> str:
    lines = ["📚 Ваши консультации:\n"]

    for index, consultation in enumerate(consultations, start=1):
        title = consultation.title or consultation.category
        lines.append(
            f"{index}. {to_html(title)} — {consultation.updated_at.astimezone(tz).strftime('%d.%m.%Y')}"
        )

    return "\n".join(lines)


def build_consultation_text(consultation, messages: list) -> str:
    lines = [
        f"📚 Консультация #{consultation.id}",
        "",
        f"Категория: {to_html(consultation.category)}",
        f"Регион: {to_html(consultation.region or 'не указан')}",
        f"Статус: {to_html(consultation.status)}",
        "",
    ]

    if not messages:
        lines.append("Сообщений пока нет.")
        return "\n".join(lines)

    for message in messages[-6:]:
        if message.role == "user":
            title = "👤 Вопрос"
        elif message.role == "assistant":
            title = "🤖 Ответ"
        else:
            title = "ℹ️ Сообщение"

        lines.append(title)
        if message.role == "assistant":
            lines.append(format_llm_answer_html(message.content[:1800]))
        else:
            lines.append(to_html(message.content[:1800]))
        lines.append("")

    return "\n".join(lines).strip()


def build_fallback_error_text(error_text: str) -> str:
    if "No reliable law chunks" in error_text:
        return (
            "Я не нашел в базе надежную норму по этому вопросу. "
            "Попробуйте уточнить ситуацию, добавить детали или выбрать другую категорию."
        )

    return error_text


async def show_main_menu(message: types.Message, state: FSMContext) -> None:
    await message.answer(build_start_text(), reply_markup=get_main_menu_kb())
    await state.set_state(MainStates.main)


async def process_question(
    message: types.Message,
    state: FSMContext,
    *,
    category: str,
    question: str,
    region: str,
    consultation_id: int | None = None,
) -> None:
    user = await orm.get_user(message.from_user.id)
    if user is None:
        await show_main_menu(message, state)
        return

    if consultation_id is None:
        consultations_today = await orm.count_user_consultations_since(
            user_id=message.from_user.id,
            since=get_today_start_utc(),
        )
        if consultations_today >= DAILY_CONSULTATION_LIMIT:
            await message.answer(
                "На сегодня достигнут лимит: 5 новых консультаций. Попробуйте завтра.",
                reply_markup=get_main_menu_kb(),
            )
            await state.set_state(MainStates.main)
            return
    else:
        message_count = await orm.count_user_messages_in_consultation(consultation_id)
        if message_count >= CONSULTATION_MESSAGE_LIMIT:
            await message.answer(
                "В этой консультации уже достигнут лимит: 5 сообщений пользователя. "
                "Создайте новую консультацию.",
                reply_markup=get_main_menu_kb(),
            )
            await state.set_state(MainStates.main)
            return

    history = []
    if consultation_id is not None:
        history = await orm.get_consultation_messages(consultation_id)

    await state.set_state(AskQuestionStates.processing)
    await message.answer("⏳ Ищу подходящие нормы закона и готовлю ответ...")

    try:
        response = await ask_rag_answer(
            user_id=message.from_user.id,
            consultation_id=consultation_id,
            question=question,
            category=category,
            region=region,
            user_type=user.user_type,
            history=history,
        )
    except RuntimeError as exc:
        await message.answer(
            build_fallback_error_text(str(exc)),
            reply_markup=get_main_menu_kb(),
        )
        await state.set_state(MainStates.main)
        return

    if user.region != region:
        await orm.set_user_region(message.from_user.id, region)

    raw_answer_text = response.get("answer", "Не удалось сформировать ответ.")
    answer_chunks = [
        format_llm_answer_html(chunk)
        for chunk in split_plain_text_chunks(raw_answer_text)
    ]
    saved_consultation_id = response.get("consultation_id")

    for index, answer_chunk in enumerate(answer_chunks):
        reply_markup = None
        if index == len(answer_chunks) - 1 and saved_consultation_id is not None:
            reply_markup = get_consultation_actions_ikb(saved_consultation_id)
        await message.answer(answer_chunk, reply_markup=reply_markup)
    await message.answer(
        "Можете задать новый вопрос или открыть историю консультаций.",
        reply_markup=get_main_menu_kb(),
    )
    await state.set_state(MainStates.main)


@client_router.message(F.text == "↩️ В меню", StateFilter("*"))
async def return_to_main_menu(message: types.Message, state: FSMContext):
    await show_main_menu(message, state)


@client_router.message(F.text == "ℹ️ Помощь", StateFilter("*"))
async def show_help(message: types.Message, state: FSMContext):
    await message.answer(build_help_text(), reply_markup=get_main_menu_kb())
    await state.set_state(MainStates.main)


@client_router.message(F.text == "⚖️ Задать вопрос", StateFilter("*"))
async def ask_question_entry(message: types.Message, state: FSMContext):
    await message.answer(
        "Выберите категорию вопроса:",
        reply_markup=get_categories_kb(),
    )
    await state.set_state(AskQuestionStates.choosing_category)


@client_router.message(
    F.text.in_(list(CATEGORY_BY_TEXT.keys())),
    AskQuestionStates.choosing_category,
)
async def set_question_category(message: types.Message, state: FSMContext):
    await state.update_data(
        category_text=message.text,
    )
    await message.answer(
        "Опишите ситуацию одним сообщением.\n\n"
        "Постарайтесь указать:\n"
        "— что произошло;\n"
        "— когда произошло;\n"
        "— с кем спор;\n"
        "— чего вы хотите добиться.",
        reply_markup=get_back_to_menu_kb(),
    )
    await state.set_state(AskQuestionStates.waiting_question)


@client_router.message(AskQuestionStates.waiting_question)
async def set_question_text(message: types.Message, state: FSMContext):
    question = (message.text or "").strip()
    if not question:
        await message.answer("Опишите ситуацию текстом одним сообщением.")
        return

    user = await orm.get_user(message.from_user.id)

    await state.update_data(question=question)

    if user and user.region:
        await message.answer(
            "Укажите регион РФ, где произошла ситуация.\n\n"
            f"Сейчас в профиле указан регион: {to_html(user.region)}\n"
            "Если он подходит, отправьте символ: -",
            reply_markup=get_back_to_menu_kb(),
        )
    else:
        await message.answer(
            "Укажите регион РФ, где произошла ситуация.\n\n"
            "Например:\n"
            "Москва\n"
            "Санкт-Петербург\n"
            "Краснодарский край\n"
            "Республика Татарстан",
            reply_markup=get_back_to_menu_kb(),
        )

    await state.set_state(AskQuestionStates.waiting_region)


@client_router.message(AskQuestionStates.waiting_region)
async def set_region_and_process(message: types.Message, state: FSMContext):
    payload = await state.get_data()
    user = await orm.get_user(message.from_user.id)
    region = (message.text or "").strip()

    if region == "-" and user and user.region:
        region = user.region

    if not region:
        await message.answer("Укажите регион РФ текстом.")
        return

    await process_question(
        message,
        state,
        category=payload["category_text"],
        question=payload["question"],
        region=region,
    )


@client_router.message(F.text == "📚 Мои консультации", StateFilter("*"))
async def show_consultations(message: types.Message, state: FSMContext):
    consultations = await orm.list_user_consultations(message.from_user.id)

    if not consultations:
        await message.answer(
            "У вас пока нет консультаций.",
            reply_markup=get_main_menu_kb(),
        )
        await state.set_state(MainStates.main)
        return

    await message.answer(
        build_consultations_text(consultations),
        reply_markup=get_consultations_ikb(consultations),
    )
    await state.set_state(MainStates.main)


@client_router.callback_query(F.data.startswith("consultation:open:"))
async def open_consultation(callback: types.CallbackQuery, state: FSMContext):
    consultation_id = int(callback.data.rsplit(":", 1)[-1])
    consultation = await orm.get_consultation(consultation_id, callback.from_user.id)

    if consultation is None:
        await callback.answer("Консультация не найдена.", show_alert=True)
        return

    messages = await orm.get_consultation_messages(consultation_id)
    await callback.message.answer(
        build_consultation_text(consultation, messages),
        reply_markup=get_consultation_actions_ikb(consultation_id),
    )
    await callback.answer()
    await state.set_state(MainStates.main)


@client_router.callback_query(F.data.startswith("consultation:continue:"))
async def continue_consultation(callback: types.CallbackQuery, state: FSMContext):
    consultation_id = int(callback.data.rsplit(":", 1)[-1])
    consultation = await orm.get_consultation(consultation_id, callback.from_user.id)

    if consultation is None:
        await callback.answer("Консультация не найдена.", show_alert=True)
        return

    await state.update_data(
        consultation_id=consultation.id,
        category_text=consultation.category,
        region=consultation.region,
    )
    await callback.message.answer(
        "Напишите следующий вопрос по этой консультации одним сообщением.",
        reply_markup=get_back_to_menu_kb(),
    )
    await callback.answer()
    await state.set_state(AskQuestionStates.waiting_continue_question)


@client_router.callback_query(F.data.startswith("consultation:delete:"))
async def delete_consultation(callback: types.CallbackQuery, state: FSMContext):
    consultation_id = int(callback.data.rsplit(":", 1)[-1])
    await orm.delete_consultation(consultation_id, callback.from_user.id)
    consultations = await orm.list_user_consultations(callback.from_user.id)

    await callback.answer("Консультация удалена.")

    if not consultations:
        await callback.message.answer(
            "Консультаций больше не осталось.",
            reply_markup=get_main_menu_kb(),
        )
        await state.set_state(MainStates.main)
        return

    await callback.message.answer(
        build_consultations_text(consultations),
        reply_markup=get_consultations_ikb(consultations),
    )
    await state.set_state(MainStates.main)


@client_router.message(AskQuestionStates.waiting_continue_question)
async def process_continue_question(message: types.Message, state: FSMContext):
    question = (message.text or "").strip()
    if not question:
        await message.answer("Напишите вопрос текстом одним сообщением.")
        return

    payload = await state.get_data()
    consultation_id = payload.get("consultation_id")
    consultation = await orm.get_consultation(consultation_id, message.from_user.id)

    if consultation is None:
        await message.answer(
            "Консультация не найдена. Начните новую.",
            reply_markup=get_main_menu_kb(),
        )
        await state.set_state(MainStates.main)
        return

    await process_question(
        message,
        state,
        category=consultation.category,
        question=question,
        region=consultation.region or payload.get("region"),
        consultation_id=consultation.id,
    )


@client_router.message(F.text == "👤 Профиль", StateFilter("*"))
async def show_profile(message: types.Message, state: FSMContext):
    user = await orm.get_user(message.from_user.id)

    if user is None:
        await show_main_menu(message, state)
        return

    await message.answer(build_profile_text(user), reply_markup=get_profile_kb())
    await state.set_state(MainStates.main)


@client_router.message(F.text == "🌍 Изменить регион", StateFilter("*"))
async def change_region_start(message: types.Message, state: FSMContext):
    await message.answer(
        "Напишите новый регион РФ одним сообщением.",
        reply_markup=get_back_to_menu_kb(),
    )
    await state.set_state(ProfileStates.waiting_region)


@client_router.message(ProfileStates.waiting_region)
async def change_region_finish(message: types.Message, state: FSMContext):
    region = (message.text or "").strip()
    if not region:
        await message.answer("Регион не должен быть пустым.")
        return

    await orm.set_user_region(message.from_user.id, region)
    user = await orm.get_user(message.from_user.id)
    await message.answer(build_profile_text(user), reply_markup=get_profile_kb())
    await state.set_state(MainStates.main)


@client_router.message(F.text == "👔 Изменить тип пользователя", StateFilter("*"))
async def change_user_type_start(message: types.Message, state: FSMContext):
    await message.answer(
        "Выберите тип пользователя:",
        reply_markup=get_user_types_kb(),
    )
    await state.set_state(ProfileStates.waiting_user_type)


@client_router.message(
    F.text.in_(list(USER_TYPE_BY_TEXT.keys())),
    ProfileStates.waiting_user_type,
)
async def change_user_type_finish(message: types.Message, state: FSMContext):
    await orm.set_user_type(message.from_user.id, USER_TYPE_BY_TEXT[message.text])
    user = await orm.get_user(message.from_user.id)
    await message.answer(build_profile_text(user), reply_markup=get_profile_kb())
    await state.set_state(MainStates.main)


@client_router.message(F.text == "🗑 Удалить мои данные", StateFilter("*"))
async def delete_profile_start(message: types.Message, state: FSMContext):
    await message.answer(
        "Удалить историю консультаций и сбросить профиль?",
        reply_markup=get_delete_profile_kb(),
    )
    await state.set_state(ProfileStates.confirm_delete)


@client_router.message(F.text == "✅ Да, удалить", ProfileStates.confirm_delete)
async def delete_profile_finish(message: types.Message, state: FSMContext):
    await orm.delete_user(message.from_user.id)
    await orm.create_user(
        user_id=message.from_user.id,
        username="@" + message.from_user.username if message.from_user.username else None,
        fullname=to_html(message.from_user.full_name),
        register_date=datetime.now(timezone.utc),
    )
    await message.answer(
        "Данные удалены. История очищена, профиль сброшен.",
        reply_markup=get_main_menu_kb(),
    )
    await state.set_state(MainStates.main)


@client_router.message(MainStates.main)
async def main_menu_fallback(message: types.Message):
    await message.answer(
        "Выберите действие в меню ниже.",
        reply_markup=get_main_menu_kb(),
    )