bot.py

import random
import asyncio
from datetime import datetime, timedelta
from typing import Optional
import uuid

from aiogram import Bot, Dispatcher, F
from aiogram.types import Message, CallbackQuery, InlineKeyboardMarkup, InlineKeyboardButton
from aiogram.filters import Command, CommandObject
from aiogram.enums import ParseMode
from aiogram.exceptions import TelegramBadRequest

from config import BOT_TOKEN, ADMIN_IDS, DEFAULT_REGISTRATION_HOURS
from database import db, Queue, Participant, QueuePhase, ParticipantStatus

bot = Bot(token=BOT_TOKEN)
dp = Dispatcher()

close_tasks: dict[str, asyncio.Task] = {}


# ========== УТИЛИТЫ ==========

def generate_queue_id() -> str:
    return uuid.uuid4().hex[:8]


def mention_user(p: Participant) -> str:
    if p.username:
        return f"@{p.username}"
    return f"[{p.full_name}](tg://user?id={p.user_id})"


def get_keyboard(queue: Queue, is_admin: bool = False) -> InlineKeyboardMarkup:
    buttons = []
    
    if queue.phase in (QueuePhase.REGISTRATION, QueuePhase.ADDITIONAL_REGISTRATION):
        buttons.append([
            InlineKeyboardButton(text="✋ Записаться", callback_data=f"join:{queue.id}"),
            InlineKeyboardButton(text="🚪 Выйти", callback_data=f"leave:{queue.id}")
        ])
    
    if queue.phase == QueuePhase.ACTIVE:
        buttons.append([
            InlineKeyboardButton(text="❌ Отказаться от места", callback_data=f"decline:{queue.id}")
        ])
    
    buttons.append([
        InlineKeyboardButton(text="📋 Показать список", callback_data=f"list:{queue.id}")
    ])
    
    if is_admin and queue.phase == QueuePhase.ACTIVE:
        buttons.append([
            InlineKeyboardButton(text="🔄 Донабор", callback_data=f"additional:{queue.id}"),
            InlineKeyboardButton(text="🔒 Закрыть", callback_data=f"close_final:{queue.id}")
        ])
    
    return InlineKeyboardMarkup(inline_keyboard=buttons)


def format_message(queue: Queue, show_list: bool = False) -> str:
    phase_text = {
        QueuePhase.REGISTRATION: "🟢 Идёт запись",
        QueuePhase.ACTIVE: "🟡 Запись закрыта",
        QueuePhase.ADDITIONAL_REGISTRATION: "🟠 Донабор",
        QueuePhase.CLOSED: "🔴 Закрыта"
    }
    
    if queue.phase in (QueuePhase.REGISTRATION, QueuePhase.ADDITIONAL_REGISTRATION):
        # Режим записи — показываем записавшихся
        participants = [p for p in db.get_participants(queue.id) if p.status != ParticipantStatus.DECLINED]
        count = len(participants)
        
        text = (
            f"📋 **Очередь: {queue.name}**\n\n"
            f"Статус: {phase_text[queue.phase]}\n"
            f"Мест: {queue.main_slots} + {queue.reserve_slots} резерв\n"
            f"Записалось: {count} чел.\n"
            f"⏰ Закрытие: {queue.closes_at.strftime('%d.%m %H:%M')}"
        )
        
        if show_list and participants:
            text += "\n\n**Записались:**\n"
            for i, p in enumerate(participants, 1):
                text += f"{i}. {p.full_name}\n"
    
    else:
        # Режим результатов
        main_list = db.get_main_participants(queue.id)
        reserve_list = db.get_reserve_participants(queue.id)
        
        text = f"🎲 **Очередь: {queue.name}**\n\n"
        text += f"Статус: {phase_text[queue.phase]}\n\n"
        
        text += "**📌 Основной список:**\n"
        if main_list:
            for i, p in enumerate(main_list, 1):
                icon = " ✅" if p.status == ParticipantStatus.PROMOTED else ""
                text += f"{i}. {mention_user(p)}{icon}\n"
        else:
            text += "(пусто)\n"
        
        if queue.reserve_slots > 0:
            text += "\n**🔄 Резерв:**\n"
            if reserve_list:
                for i, p in enumerate(reserve_list, 1):
                    text += f"{i}. {mention_user(p)}\n"
            else:
                text += "(пусто)\n"
        
        free_main = queue.main_slots - len(main_list)
        if free_main > 0 and queue.phase != QueuePhase.CLOSED:
            text += f"\n⚠️ **Свободно мест: {free_main}**"
    
    return text


async def update_message(queue: Queue, show_list: bool = False):
    try:
        text = format_message(queue, show_list)
        keyboard = get_keyboard(queue, is_admin=True)
        
        await bot.edit_message_text(
            chat_id=queue.chat_id,
            message_id=queue.message_id,
            text=text,
            reply_markup=keyboard if queue.phase != QueuePhase.CLOSED else None,
            parse_mode=ParseMode.MARKDOWN
        )
    except TelegramBadRequest:
        pass


# ========== СОЗДАНИЕ ОЧЕРЕДИ ==========

@dp.message(Command("create"))
async def create_queue(message: Message, command: CommandObject):
    chat_id = message.chat.id
    user_id = message.from_user.id
    
    if user_id not in ADMIN_IDS:
        return await message.answer("⛔ Только админ может создавать очереди")
    
    args = command.args.split() if command.args else []
    
    if len(args) < 3:
        return await message.answer(
            "❌ **Формат:**\n"
            "`/create <название> <мест> <резерв> [часов]`\n\n"
            "**Пример:** `/create Матан 5 2`",
            parse_mode=ParseMode.MARKDOWN
        )
    
    try:
        name = args[0]
        main_slots = int(args[1])
        reserve_slots = int(args[2])
        hours = float(args[3]) if len(args) > 3 else DEFAULT_REGISTRATION_HOURS
    except ValueError:
        return await message.answer("❌ Неверные параметры")
    
    if main_slots < 1 or reserve_slots < 0 or hours <= 0:
        return await message.answer("❌ Некорректные значения")
    
    existing = db.find_queue_by_name(chat_id, name)
    if existing:
        return await message.answer(f"❌ Очередь '{name}' уже существует")
    
    now = datetime.now()
    queue = Queue(
        id=generate_queue_id(),
        name=name,
        chat_id=chat_id,
        message_id=None,
        main_slots=main_slots,
        reserve_slots=reserve_slots,
        phase=QueuePhase.REGISTRATION,
        created_at=now,
        closes_at=now + timedelta(hours=hours)
    )
    
    db.create_queue(queue)
    
    msg = await message.answer(
        format_message(queue),
        reply_markup=get_keyboard(queue),
        parse_mode=ParseMode.MARKDOWN
    )
    
    queue.message_id = msg.message_id
    db.update_queue(queue)
    
    schedule_auto_close(queue)
    
    try:
        await message.delete()
    except:
        pass


def schedule_auto_close(queue: Queue):
    delay = (queue.closes_at - datetime.now()).total_seconds()
    if delay > 0:
        task = asyncio.create_task(auto_close_queue(queue.id, delay))
        close_tasks[queue.id] = task


async def auto_close_queue(queue_id: str, delay: float):
    await asyncio.sleep(delay)
    queue = db.get_queue(queue_id)
    if queue and queue.phase == QueuePhase.REGISTRATION:
        await finalize_registration(queue)
    elif queue and queue.phase == QueuePhase.ADDITIONAL_REGISTRATION:
        await finalize_additional(queue)


# ========== ЗАПИСЬ ==========

@dp.callback_query(F.data.startswith("join:"))
async def join_queue(callback: CallbackQuery):
    queue_id = callback.data.split(":")[1]
    queue = db.get_queue(queue_id)
    
    if not queue:
        return await callback.answer("❌ Очередь не найдена", show_alert=True)
    
    if queue.phase not in (QueuePhase.REGISTRATION, QueuePhase.ADDITIONAL_REGISTRATION):
        return await callback.answer("❌ Запись закрыта", show_alert=True)
    
    user = callback.from_user
    existing = db.get_participant(queue_id, user.id)
    
    if existing and existing.status != ParticipantStatus.DECLINED:
        return await callback.answer("✅ Вы уже записаны!", show_alert=True)
    
    participant = Participant(
        id=0,
        queue_id=queue_id,
        user_id=user.id,
        username=user.username,
        full_name=user.full_name,
        position=None,
        status=ParticipantStatus.PENDING,
        joined_at=datetime.now()
    )
    
    db.add_participant(participant)
    await update_message(queue)
    
    count = db.count_participants(queue_id)
    await callback.answer(f"✅ Вы записаны! Всего: {count}", show_alert=True)


@dp.callback_query(F.data.startswith("leave:"))
async def leave_queue(callback: CallbackQuery):
    queue_id = callback.data.split(":")[1]
    queue = db.get_queue(queue_id)
    
    if not queue:
        return await callback.answer("❌ Очередь не найдена", show_alert=True)
    
    if queue.phase not in (QueuePhase.REGISTRATION, QueuePhase.ADDITIONAL_REGISTRATION):
        return await callback.answer("❌ Используйте 'Отказаться от места'", show_alert=True)
    
    if not db.remove_participant(queue_id, callback.from_user.id):
        return await callback.answer("❌ Вы не записаны", show_alert=True)
    
    await update_message(queue)
    await callback.answer("🚪 Вы вышли из очереди", show_alert=True)


@dp.callback_query(F.data.startswith("list:"))
async def show_list(callback: CallbackQuery):
    queue_id = callback.data.split(":")[1]
    queue = db.get_queue(queue_id)
    
    if not queue:
        return await callback.answer("❌ Очередь не найдена", show_alert=True)
    
    await update_message(queue, show_list=True)
    await callback.answer()


# ========== ЗАКРЫТИЕ ЗАПИСИ ==========

@dp.message(Command("close"))
async def close_queue_command(message: Message, command: CommandObject):
    chat_id = message.chat.id
    user_id = message.from_user.id
    
    if user_id not in ADMIN_IDS:
        return await message.answer("⛔ Только админ")
    
    if not command.args:
        queues_list = db.get_active_queues(chat_id)
        if not queues_list:
            return await message.answer("❌ Нет активных очередей")
        
        lines = ["**Активные очереди:**"]
        for q in queues_list:
            lines.append(f"• `{q.name}`")
        return await message.answer("\n".join(lines), parse_mode=ParseMode.MARKDOWN)
    
    name = command.args.strip()
    queue = db.find_queue_by_name(chat_id, name) or db.get_queue(name)
    
    if not queue or queue.chat_id != chat_id:
        return await message.answer("❌ Очередь не найдена")
    
    if queue.phase == QueuePhase.REGISTRATION:
        await finalize_registration(queue)
    elif queue.phase == QueuePhase.ADDITIONAL_REGISTRATION:
        await finalize_additional(queue)
    else:
        return await message.answer("❌ Очередь уже закрыта для записи")
    
    try:
        await message.delete()
    except:
        pass


async def finalize_registration(queue: Queue):
    if queue.id in close_tasks:
        close_tasks[queue.id].cancel()
        del close_tasks[queue.id]
    
    participants = db.get_pending_participants(queue.id)
    random.shuffle(participants)
    
    for i, p in enumerate(participants):
        if i < queue.main_slots:
            status = ParticipantStatus.MAIN
        elif i < queue.main_slots + queue.reserve_slots:
            status = ParticipantStatus.RESERVE
        else:
            status = ParticipantStatus.WAITING
        
        db.update_participant_status(queue.id, p.user_id, status, position=i + 1)
    
    queue.phase = QueuePhase.ACTIVE
    db.update_queue(queue)
    
    await update_message(queue)


# ========== ОТКАЗ ОТ МЕСТА ==========

@dp.callback_query(F.data.startswith("decline:"))
async def decline_spot(callback: CallbackQuery):
    queue_id = callback.data.split(":")[1]
    queue = db.get_queue(queue_id)
    user_id = callback.from_user.id
    
    if not queue:
        return await callback.answer("❌ Очередь не найдена", show_alert=True)
    
    if queue.phase not in (QueuePhase.ACTIVE, QueuePhase.ADDITIONAL_REGISTRATION):
        return await callback.answer("❌ Нельзя отказаться сейчас", show_alert=True)
    
    participant = db.get_participant(queue_id, user_id)
    
    if not participant:
        return await callback.answer("❌ Вы не в очереди", show_alert=True)
    
    if participant.status in (ParticipantStatus.DECLINED, ParticipantStatus.WAITING, ParticipantStatus.PENDING):
        return await callback.answer("❌ Вы не в основном списке/резерве", show_alert=True)
    
    was_main = participant.status in (ParticipantStatus.MAIN, ParticipantStatus.PROMOTED)
    
    db.update_participant_status(queue_id, user_id, ParticipantStatus.DECLINED)
    
    await callback.answer("✅ Вы отказались от места", show_alert=True)
    
    if was_main:
        await promote_reserve(queue)
    
    await update_message(queue)


async def promote_reserve(queue: Queue):
    reserve = db.get_first_reserve(queue.id)
    if not reserve:
        return
    
    main_list = db.get_main_participants(queue.id)
    new_position = len(main_list) + 1
    
    db.update_participant_status(queue.id, reserve.user_id, ParticipantStatus.PROMOTED, position=new_position)
    
    # Уведомление резервисту
    try:
        await bot.send_message(
            reserve.user_id,
            f"🎉 Вы продвинуты в основной список очереди **{queue.name}**!",
            parse_mode=ParseMode.MARKDOWN
        )
    except:
        pass


# ========== ДОНАБОР ==========

@dp.callback_query(F.data.startswith("additional:"))
async def start_additional(callback: CallbackQuery):
    if callback.from_user.id not in ADMIN_IDS:
        return await callback.answer("⛔ Только админ", show_alert=True)
    
    queue_id = callback.data.split(":")[1]
    queue = db.get_queue(queue_id)
    
    if not queue or queue.phase != QueuePhase.ACTIVE:
        return await callback.answer("❌ Донабор недоступен", show_alert=True)
    
    main_list = db.get_main_participants(queue.id)
    free_slots = queue.main_slots - len(main_list)
    
    if free_slots <= 0:
        return await callback.answer("✅ Все места заняты", show_alert=True)
    
    queue.phase = QueuePhase.ADDITIONAL_REGISTRATION
    queue.closes_at = datetime.now() + timedelta(hours=DEFAULT_REGISTRATION_HOURS)
    db.update_queue(queue)
    
    schedule_auto_close(queue)
    
    await update_message(queue)
    await callback.answer(f"🔄 Донабор открыт на {free_slots} мест", show_alert=True)


async def finalize_additional(queue: Queue):
    if queue.id in close_tasks:
        close_tasks[queue.id].cancel()
        del close_tasks[queue.id]
    
    pending = db.get_pending_participants(queue.id)
    random.shuffle(pending)
    
    main_list = db.get_main_participants(queue.id)
    free_slots = queue.main_slots - len(main_list)
    
    for i, p in enumerate(pending):
        if i < free_slots:
            db.update_participant_status(queue.id, p.user_id, ParticipantStatus.PROMOTED, position=len(main_list) + i + 1)
        else:
            db.update_participant_status(queue.id, p.user_id, ParticipantStatus.WAITING)
    
    queue.phase = QueuePhase.ACTIVE
    db.update_queue(queue)
    
    await update_message(queue)


# ========== ЗАКРЫТИЕ ПОЛНОЕ ==========

@dp.callback_query(F.data.startswith("close_final:"))
async def close_final(callback: CallbackQuery):
    if callback.from_user.id not in ADMIN_IDS:
        return await callback.answer("⛔ Только админ", show_alert=True)
    
    queue_id = callback.data.split(":")[1]
    queue = db.get_queue(queue_id)
    
    if not queue:
        return await callback.answer("❌ Не найдена", show_alert=True)
    
    queue.phase = QueuePhase.CLOSED
    db.update_queue(queue)
    
    await update_message(queue)
    await callback.answer("🔒 Очередь закрыта", show_alert=True)


# ========== СЛУЖЕБНЫЕ ==========

@dp.message(Command("queues"))
async def list_queues(message: Message):
    queues_list = db.get_active_queues(message.chat.id)
    
    if not queues_list:
        return await message.answer("📋 Нет активных очередей")
    
    lines = ["**📋 Активные очереди:**"]
    for q in queues_list:
        count = db.count_participants(q.id)
        lines.append(f"• **{q.name}** — {count} чел.")
    
    await message.answer("\n".join(lines), parse_mode=ParseMode.MARKDOWN)


@dp.message(Command("help", "start"))
async def help_command(message: Message):
    text = "**🎲 Бот для очередей**\n\n"
    
    if message.from_user.id in ADMIN_IDS:
        text += (
            "`/create <имя> <мест> <резерв> [часов]`\n"
            "`/close <имя>` — закрыть досрочно\n"
            "`/queues` — список очередей\n\n"
            "**Пример:** `/create Матан 5 2`"
        )
    else:
        text += "Нажмите кнопку под сообщением очереди чтобы записаться."
    
    await message.answer(text, parse_mode=ParseMode.MARKDOWN)


async def restore_timers():
    for queue in db.get_registration_queues():
        delay = (queue.closes_at - datetime.now()).total_seconds()
        if delay > 0:
            schedule_auto_close(queue)
        else:
            if queue.phase == QueuePhase.REGISTRATION:
                await finalize_registration(queue)
            elif queue.phase == QueuePhase.ADDITIONAL_REGISTRATION:
                await finalize_additional(queue)


async def main():
    print("Запуск бота...")
    await restore_timers()
    await dp.start_polling(bot)


if __name__ == "__main__":
    asyncio.run(main())