Click/main.py

"""
Главный файл бота с сохранением задач.
"""

import asyncio
import logging
from types import SimpleNamespace
from aiogram import Bot, Dispatcher
from aiogram.types import BotCommand
from aiogram.client.session.aiohttp import AiohttpSession
from config.settings import settings
from core.logger import setup_logger
from managers.proxy_manager import ProxyManager
from managers.background_tasks import BackgroundTaskManager
from services.browser_service import BrowserService
from handlers.commands import BotInterface
from managers.storage import TaskStorage, ChatStorage

logger = setup_logger(__name__)


async def set_bot_commands(bot: Bot, is_admin: bool = False):
    """Устанавливает команды бота в меню."""
    if is_admin:
        commands = [
            BotCommand(command="start", description="🏠 Главное меню"),
            BotCommand(command="logout", description="🔒 Выйти"),
            BotCommand(command="streamers", description="📺 Стримеры"),
            BotCommand(command="tasks", description="📊 Задачи"),
            BotCommand(command="status", description="📈 Статус"),
        ]
    else:
        commands = [
            BotCommand(command="start", description="🤖 Старт"),
        ]
    await bot.set_my_commands(commands)


class BotApplication:
    
    def __init__(self):
        self.bot: Bot = None
        self.dispatcher: Dispatcher = None
        
        self.proxy_manager = ProxyManager()
        self.background_tasks = BackgroundTaskManager()
        self.browser_service = BrowserService(self.proxy_manager)
        self.storage = TaskStorage()
        self.chat_storage = ChatStorage()
        
        self.interface = BotInterface(
            background_tasks=self.background_tasks,
            proxy_manager=self.proxy_manager,
            browser_service=self.browser_service,
            storage=self.storage,
            chat_storage=self.chat_storage,
            bot_ref=self,
        )
    
    async def _update_bot_commands(self):
        is_admin = any(
            self.interface.auth_manager.is_admin(uid)
            for uid in self.interface.auth_manager._sessions
        )
        await set_bot_commands(self.bot, is_admin=is_admin)
    
    async def initialize(self):
        """Инициализация бота."""
        logger.info("Initializing bot...")
        
        if not settings.BOT_TOKEN:
            raise ValueError("BOT_TOKEN not set!")
         

        session = AiohttpSession(proxy=settings.TELEGRAM_PROXY) if settings.TELEGRAM_PROXY else AiohttpSession()
        self.bot = Bot(token=settings.BOT_TOKEN, session=session)
        self.dispatcher = Dispatcher()
        
        self.interface.register(self.dispatcher)
        await self._update_bot_commands()

        # Восстанавливаем задачи
        await self._restore_tasks()
        
        logger.info(f"Proxies: {self.proxy_manager.count}")
        logger.info("Bot initialized!")
    
    async def _restore_tasks(self):
        """Восстанавливает задачи после перезапуска."""
        saved_tasks = await self.storage.load_tasks()
        
        if not saved_tasks:
            logger.info("No tasks to restore")
            return
        
        restored = 0
        notify_users: dict = {}  # user_id -> chat_id

        for task_id, params in saved_tasks.items():
            # Завершённые и остановленные — загружаем в память но не запускаем
            if params.completed or params.stopped:
                await self.interface.task_manager.add_task(task_id, params)
                continue

            # Активные задачи ставим на паузу и перезапускаем
            params.paused = True
            await self.interface.task_manager.add_task(task_id, params)

            if params.task_type == "twitch_irc" and params.chat_id:
                msg_mock = SimpleNamespace(
                    chat=SimpleNamespace(id=params.chat_id),
                    bot=self.bot,
                )
                await self.background_tasks.start_task(
                    task_id=task_id,
                    coro=self.interface._run_twitch(task_id, params, msg_mock),
                    task_type="twitch_irc",
                    metadata={'type': 'twitch_irc', 'channel': params.channel, 'chat_id': params.chat_id}
                )
                restored += 1
            elif params.task_type == "user_visit" and params.chat_id:
                await self.background_tasks.start_task(
                    task_id=task_id,
                    coro=self.interface._run_url_visits(task_id, params, self.bot),
                    task_type="user_visit",
                    metadata={'type': 'user_visit', 'url': params.url, 'chat_id': params.chat_id}
                )
                restored += 1

            if params.user_id and params.chat_id:
                notify_users[params.user_id] = params.chat_id

        # Уведомляем пользователей о рестарте
        for uid, chat_id in notify_users.items():
            try:
                await self.bot.send_message(
                    chat_id,
                    "🔄 Бот был перезапущен\n\n"
                    "Ваши задачи приостановлены.\n"
                    "Нажмите ▶️ в списке задач для возобновления."
                )
            except Exception as e:
                logger.warning(f"Failed to notify user {uid}: {e}")
        
        logger.info(f"🔄 Restored {restored} tasks")
    
    async def start(self):
        """Запуск бота."""
        await self.initialize()
        logger.info("Starting polling...")
        await self.dispatcher.start_polling(self.bot)
    
    async def stop(self):
        """Остановка с сохранением."""
        logger.info("Shutting down...")
        
        # Сохраняем задачи
        tasks = self.interface.task_manager._tasks
        await self.storage.save_tasks(tasks)
        
        await self.background_tasks.cancel_all()
        await self.browser_service.cleanup()
        if self.bot:
            await self.bot.session.close()
        logger.info("Bot stopped")


async def main():
    app = BotApplication()
    try:
        await app.start()
    except KeyboardInterrupt:
        logger.info("Interrupted")
    finally:
        await app.stop()


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