main.py

import asyncio
import logging
import re
import time

from aiogram import Bot, Dispatcher, F, Router
from aiogram.client.default import DefaultBotProperties
from aiogram.enums import ParseMode
from aiogram.filters import Command, CommandStart
from aiogram.types import BotCommand, CallbackQuery, Message

import config
import session_manager as sm
import ui
from flavors import FLAVORS, flood_seconds

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s")
log = logging.getLogger(__name__)

router = Router()

_HASH_RE = re.compile(r"^[0-9a-fA-F]{32}$")

async def panel(bot: Bot, chat_id: int, text: str, kb=None, effect: str | None = None) -> None:
    old_id = sm.panels.get(chat_id)

    if effect:
        if old_id:
            try:
                await bot.delete_message(chat_id, old_id)
            except Exception:
                pass
        msg = await bot.send_message(chat_id, text, reply_markup=kb, message_effect_id=effect)
        sm.panels[chat_id] = msg.message_id
        return

    if old_id:
        try:
            await bot.edit_message_text(text, chat_id=chat_id, message_id=old_id, reply_markup=kb)
            return
        except Exception:
            pass

    msg = await bot.send_message(chat_id, text, reply_markup=kb)
    sm.panels[chat_id] = msg.message_id


async def _delete_input(message: Message) -> None:
    try:
        await message.delete()
    except Exception:
        pass

# ── Entry points
@router.message(CommandStart())
async def cmd_start(message: Message, bot: Bot) -> None:
    await sm.drop(message.from_user.id, message.chat.id, bot)
    me = await bot.get_me()
    await panel(bot, message.chat.id, ui.welcome(me.username), ui.kb_welcome(), effect=config.EFFECT_WELCOME)


@router.message(Command("ping"))
async def cmd_ping(message: Message, bot: Bot) -> None:
    t0 = time.monotonic()
    sent = await message.answer("🏓")
    dt_ms = (time.monotonic() - t0) * 1000
    await sent.edit_text(ui.pong(dt_ms))


@router.callback_query(F.data == "gen_start")
async def cb_gen_start(cq: CallbackQuery, bot: Bot) -> None:
    await sm.drop(cq.from_user.id, cq.message.chat.id, bot)
    sm.set_step(cq.from_user.id, step="library")
    sm.start_timer(cq.from_user.id, cq.message.chat.id, bot)
    await cq.answer()
    await panel(bot, cq.message.chat.id, ui.step_library(), ui.kb_library())


@router.callback_query(F.data.startswith("lib:"))
async def cb_choose_library(cq: CallbackQuery, bot: Bot) -> None:
    user_id = cq.from_user.id
    state = sm.get(user_id)
    if not state:
        await cq.answer("Session expired, tap Generate to start over.", show_alert=True)
        return
    flavor_key = cq.data.split(":", 1)[1]
    sm.set_step(user_id, step="credentials", flavor=flavor_key)
    sm.start_timer(user_id, cq.message.chat.id, bot)
    await cq.answer(f"{FLAVORS[flavor_key].name} selected")
    await panel(bot, cq.message.chat.id, ui.step_credentials(), ui.kb_credentials())


@router.callback_query(F.data.startswith("cred:"))
async def cb_choose_credentials(cq: CallbackQuery, bot: Bot) -> None:
    user_id = cq.from_user.id
    state = sm.get(user_id)
    if not state:
        await cq.answer("Session expired, tap Generate to start over.", show_alert=True)
        return

    mode = cq.data.split(":", 1)[1]
    sm.start_timer(user_id, cq.message.chat.id, bot)
    await cq.answer()

    if mode == "default":
        sm.set_step(user_id, step="phone", api_id=config.API_ID, api_hash=config.API_HASH)
        await panel(bot, cq.message.chat.id, ui.ask_phone(), ui.kb_cancel_only())
    else:
        sm.set_step(user_id, step="api_id")
        await panel(bot, cq.message.chat.id, ui.ask_api_id(), ui.kb_cancel_only())


@router.callback_query(F.data == "gen_cancel")
async def cb_cancel(cq: CallbackQuery, bot: Bot) -> None:
    await sm.drop(cq.from_user.id, cq.message.chat.id, bot)
    await cq.answer("Cancelled.")
    await panel(bot, cq.message.chat.id, ui.cancelled(), ui.kb_restart())

# ── Free-text steps 
@router.message(F.text & ~F.text.startswith("/"))
async def on_text(message: Message, bot: Bot) -> None:
    user_id = message.from_user.id
    state = sm.get(user_id)
    if not state:
        return  # no active flow — ignore stray text

    step = state.get("step")
    text = message.text.strip()

    handlers = {
        "api_id": _handle_api_id,
        "api_hash": _handle_api_hash,
        "phone": _handle_phone,
        "code": _handle_code,
        "password": _handle_password,
    }
    handler = handlers.get(step)
    if handler:
        await handler(message, text, state, bot)


async def _handle_api_id(message: Message, text: str, state: dict, bot: Bot) -> None:
    await _delete_input(message)
    if not text.isdigit():
        await panel(bot, message.chat.id, ui.invalid_api_id(), ui.kb_cancel_only())
        return
    sm.set_step(message.from_user.id, step="api_hash", api_id=int(text))
    sm.start_timer(message.from_user.id, message.chat.id, bot)
    await panel(bot, message.chat.id, ui.ask_api_hash(), ui.kb_cancel_only())


async def _handle_api_hash(message: Message, text: str, state: dict, bot: Bot) -> None:
    await _delete_input(message)
    if not _HASH_RE.match(text):
        await panel(bot, message.chat.id, ui.invalid_api_hash(), ui.kb_cancel_only())
        return
    sm.set_step(message.from_user.id, step="phone", api_hash=text)
    sm.start_timer(message.from_user.id, message.chat.id, bot)
    await panel(bot, message.chat.id, ui.ask_phone(), ui.kb_cancel_only())


async def _handle_phone(message: Message, phone: str, state: dict, bot: Bot) -> None:
    await _delete_input(message)
    if not phone.startswith("+"):
        phone = "+" + phone
    if not 5 <= len(phone) <= 18:
        await panel(bot, message.chat.id, ui.invalid_phone(), ui.kb_cancel_only())
        return

    user_id = message.from_user.id
    flavor = FLAVORS[state["flavor"]]
    await panel(bot, message.chat.id, ui.sending_code(phone), ui.kb_cancel_only())

    client = flavor.client(user_id, state["api_id"], state["api_hash"])
    sm.set_step(user_id, client=client)
    try:
        await client.connect()
        phone_code_hash = await flavor.send_code(client, phone)
        sm.set_step(user_id, step="code", phone=phone, phone_code_hash=phone_code_hash)
        sm.start_timer(user_id, message.chat.id, bot)
        await panel(bot, message.chat.id, ui.ask_code(phone), ui.kb_cancel_only())
    except flavor.phone_invalid:
        await _fail(user_id, message.chat.id, bot, ui.invalid_phone())
    except flavor.flood as e:
        await _fail(user_id, message.chat.id, bot, ui.flood_wait(flood_seconds(e)))
    except Exception as e:
        log.exception("phone step failed for %s", user_id)
        await _fail(user_id, message.chat.id, bot, ui.generic_error(str(e)))


async def _handle_code(message: Message, code: str, state: dict, bot: Bot) -> None:
    await _delete_input(message)
    clean_code = code.replace(" ", "").replace(".", "").replace("-", "")
    user_id = message.from_user.id
    flavor = FLAVORS[state["flavor"]]

    try:
        await flavor.sign_in(state["client"], state["phone"], state["phone_code_hash"], clean_code)
    except flavor.need_password:
        sm.set_step(user_id, step="password")
        sm.start_timer(user_id, message.chat.id, bot)
        await panel(bot, message.chat.id, ui.ask_password(), ui.kb_cancel_only())
        return
    except flavor.code_invalid:
        await panel(bot, message.chat.id, ui.code_invalid(), ui.kb_cancel_only())
        return
    except flavor.code_expired:
        await _fail(user_id, message.chat.id, bot, ui.generic_error("That code expired. Start over with Generate."))
        return
    except flavor.flood as e:
        await _fail(user_id, message.chat.id, bot, ui.flood_wait(flood_seconds(e)))
        return
    except Exception as e:
        log.exception("code step failed for %s", user_id)
        await _fail(user_id, message.chat.id, bot, ui.generic_error(str(e)))
        return

    await _finish(user_id, message.chat.id, bot, state)


async def _handle_password(message: Message, password: str, state: dict, bot: Bot) -> None:
    await _delete_input(message)
    user_id = message.from_user.id
    flavor = FLAVORS[state["flavor"]]

    try:
        await flavor.check_password(state["client"], state["phone"], password)
    except flavor.password_invalid:
        await panel(bot, message.chat.id, ui.password_invalid(), ui.kb_cancel_only())
        return
    except Exception as e:
        log.exception("password step failed for %s", user_id)
        await _fail(user_id, message.chat.id, bot, ui.generic_error(str(e)))
        return

    await _finish(user_id, message.chat.id, bot, state)


async def _fail(user_id: int, chat_id: int, bot: Bot, text: str) -> None:
    await sm.drop(user_id, chat_id, bot)
    await panel(bot, chat_id, text, ui.kb_restart())


async def _finish(user_id: int, chat_id: int, bot: Bot, state: dict) -> None:
    flavor = FLAVORS[state["flavor"]]
    client = state["client"]
    try:
        string_session = await flavor.export(client)
    except Exception as e:
        log.exception("session export failed for %s", user_id)
        await _fail(user_id, chat_id, bot, ui.generic_error(str(e)))
        return
    finally:
        try:
            await client.disconnect()
        except Exception:
            pass

    task = state.get("timer")
    if task:
        task.cancel()
    sm.states.pop(user_id, None)
    old_notice = sm.notices.pop(chat_id, None)
    if old_notice:
        try:
            await bot.delete_message(chat_id, old_notice)
        except Exception:
            pass

    me = await bot.get_me()
    await panel(bot, chat_id, ui.success(flavor.name, string_session, me.username), effect=config.EFFECT_SUCCESS)


# ── Lifecycle ─────────────────────────────────────────────────────────

async def _on_startup(bot: Bot) -> None:
    await bot.set_my_commands([
        BotCommand(command="start", description="Open the session generator"),
        BotCommand(command="ping", description="Check bot latency"),
    ])
    log.info("Bot commands registered.")


async def main() -> None:
    bot = Bot(token=config.BOT_TOKEN, default=DefaultBotProperties(parse_mode=ParseMode.HTML))
    dp = Dispatcher()
    dp.include_router(router)
    dp.startup.register(_on_startup)

    log.info("Starting bot...")
    await bot.delete_webhook(drop_pending_updates=True)
    await dp.start_polling(bot)


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