session_manager.py

import asyncio
import logging
from aiogram import Bot
import config
import ui

log = logging.getLogger(__name__)

states: dict[int, dict] = {}

panels: dict[int, int] = {}

notices: dict[int, int] = {}

def get(user_id: int) -> dict | None:
    return states.get(user_id)

def set_step(user_id: int, **fields) -> None:
    states.setdefault(user_id, {}).update(fields)

async def _disconnect(client) -> None:
    if not client:
        return
    try:
        await client.disconnect()
    except Exception:
        pass

async def drop(user_id: int, chat_id: int, bot: Bot) -> None:
    """End the flow entirely: cancel the timer, disconnect any live client,
    forget the state, and clear the countdown notice if one is showing."""
    state = states.pop(user_id, None)
    if state:
        task = state.get("timer")
        if task:
            task.cancel()
        await _disconnect(state.get("client"))

    notice_id = notices.pop(chat_id, None)
    if notice_id:
        try:
            await bot.delete_message(chat_id, notice_id)
        except Exception:
            pass

async def _notify(bot: Bot, chat_id: int, text: str) -> None:
    old = notices.pop(chat_id, None)
    if old:
        try:
            await bot.delete_message(chat_id, old)
        except Exception:
            pass
    try:
        msg = await bot.send_message(chat_id, text, message_effect_id=config.EFFECT_NOTICE)
        notices[chat_id] = msg.message_id
    except Exception:
        log.exception("failed to send countdown notice to %s", chat_id)


async def _expire(user_id: int, chat_id: int, bot: Bot) -> None:
    state = states.pop(user_id, None)
    if state:
        await _disconnect(state.get("client"))

    old_notice = notices.pop(chat_id, None)
    if old_notice:
        try:
            await bot.delete_message(chat_id, old_notice)
        except Exception:
            pass

    panel_id = panels.get(chat_id)
    text = ui.expired()
    kb = ui.kb_restart()
    if panel_id:
        try:
            await bot.edit_message_text(text, chat_id=chat_id, message_id=panel_id, reply_markup=kb)
            return
        except Exception:
            pass
    msg = await bot.send_message(chat_id, text, reply_markup=kb)
    panels[chat_id] = msg.message_id


async def _watch(user_id: int, chat_id: int, bot: Bot) -> None:
    try:
        first_wait = config.FLOW_TIMEOUT_SECONDS - 60
        if first_wait > 0:
            await asyncio.sleep(first_wait)
            await _notify(bot, chat_id, ui.notice_60s())
        await asyncio.sleep(30)
        await _notify(bot, chat_id, ui.notice_30s())
        await asyncio.sleep(20)
        await _notify(bot, chat_id, ui.notice_10s())
        await asyncio.sleep(10)
        await _expire(user_id, chat_id, bot)
    except asyncio.CancelledError:
        pass


def start_timer(user_id: int, chat_id: int, bot: Bot) -> None:
    """(Re)start the 5-minute countdown for a user's active flow."""
    state = states.setdefault(user_id, {})
    old = state.get("timer")
    if old:
        old.cancel()
    state["timer"] = asyncio.create_task(_watch(user_id, chat_id, bot))