cyclop/telegram_helper.py

import os
import logging
import sys
import asyncio
import time
from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes, CallbackQueryHandler
from browse_command import browse_command, button_callback

class TelegramHelper:
    # --- Constants for magic strings (paths are now instance vars) ---
    CLAUDE_REBOOT_TARGET = 'claude'
    HTML_QUOTE_BLOCK_START = '<blockquote expandable><b>Thinking...</b>'
    HTML_QUOTE_BLOCK_END = '</blockquote>'
    DEFAULT_REBOOT_CLAUDE_FILE = '.reboot_claude' # Default value
    DEFAULT_REBOOT_FILE = '.doreboot'          # Default value

    def __init__(self, bot, reboot_claude_file_path: str | None = None, reboot_file_path: str | None = None): # MODIFIED
        self.bot = bot
        self.telegram_bot_token = os.getenv('TELEGRAM_BOT_TOKEN')
        self.start_time = time.time()
        # MODIFIED: Store configurable paths
        self.reboot_claude_file = reboot_claude_file_path or self.DEFAULT_REBOOT_CLAUDE_FILE
        self.reboot_file = reboot_file_path or self.DEFAULT_REBOOT_FILE

    async def start(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
        await self.bot.start()
        await update.message.reply_text(
            "Hello! I'm your AI assistant. How can I help you today?"
        )

    async def clear(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
        user_id = update.effective_user.id
        self.bot.clear_conversation_history(user_id)
        await update.message.reply_text("Conversation history cleared. Let's start fresh!")

    async def status(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
        status_message = await self.bot.get_bot_status()
        await update.message.reply_text(status_message)

    async def switch(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
        if hasattr(self.bot, 'switch_model'):
            status_message = await self.bot.switch_model()
            await update.message.reply_text(status_message)
        else:
            await update.message.reply_text("Model switching is not supported for this bot.")

    async def update_status_message(self, context: ContextTypes.DEFAULT_TYPE, chat_id: int, message_id: int, status: str):
        keyboard = [
            [InlineKeyboardButton("Abort", callback_data='abort')]
        ]
        reply_markup = InlineKeyboardMarkup(keyboard)
        await context.bot.edit_message_text(
            chat_id=chat_id,
            message_id=message_id,
            text=f"Current status: {status}",
            reply_markup=reply_markup
        )

    async def handle_message(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
        try:
            user_id = update.effective_user.id
            user_message = update.message.text

            logging.info(f"Message from user {user_id}: {user_message}")

            status_message = await update.message.reply_text("Processing your request...", reply_markup=InlineKeyboardMarkup([[InlineKeyboardButton("Abort", callback_data='abort')]]))
            self.bot.set_processing_status(user_id, status_message.message_id)

            response = await self.bot.handle_message(user_id, user_message)

            await context.bot.delete_message(chat_id=update.effective_chat.id, message_id=status_message.message_id)
            self.bot.clear_processing_status(user_id)

            response = response.replace("<think>", self.HTML_QUOTE_BLOCK_START).replace("</think>", self.HTML_QUOTE_BLOCK_END)

            if len(response) > 4096:
                chunks = [response[i:i + 4096] for i in range(0, len(response), 4096)]
                for chunk in chunks:
                    await update.message.reply_text(chunk)
                    await asyncio.sleep(0.1)
            else:
                await update.message.reply_text(response)

        except Exception as e:
            logging.error(f"An error occurred: {str(e)}")
            await update.message.reply_text("Sorry, an error occurred while processing your request.")

    async def abort_processing(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
        query = update.callback_query
        await query.answer()

        user_id = query.from_user.id
        result = await self.bot.abort_processing(user_id)
        await query.edit_message_text(text=result)

    async def reboot(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
        user_message = update.message.text.split()
        if len(user_message) > 1 and user_message[1].lower() == self.CLAUDE_REBOOT_TARGET:
            # MODIFIED: Use instance variable
            open(self.reboot_claude_file, 'w').close()

        if update:
            await update.message.reply_text("Rebooting the bot...")
        logging.info("Received reboot command. Exiting process...")
        # MODIFIED: Use instance variable
        reboot_f_path = self.reboot_file
        if not os.path.exists(reboot_f_path):
            with open(reboot_f_path, 'w') as f:
                f.write(str(update.effective_chat.id) if update else "")
        sys.exit(0)

    async def check_doreboot_file(self, application: Application):
        # MODIFIED: Use instance variable
        reboot_f_path = self.reboot_file
        if os.path.exists(reboot_f_path):
            with open(reboot_f_path, 'r') as f:
                chat_id = f.read().strip()
            if chat_id:
                await application.bot.send_message(chat_id=chat_id, text="The application has finished initializing.")
            os.remove(reboot_f_path)

    async def browse(self, update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
        await browse_command(update, context, self.bot)

    def run(self):
        application = Application.builder().token(self.telegram_bot_token).build()

        application.add_handler(CommandHandler("start", self.start))
        application.add_handler(CommandHandler("clear", self.clear))
        application.add_handler(CommandHandler("switch", self.switch))
        application.add_handler(CommandHandler("status", self.status))
        application.add_handler(CommandHandler("reboot", self.reboot))
        application.add_handler(CommandHandler("browse", self.browse))
        application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, self.handle_message))
        application.add_handler(CallbackQueryHandler(self.abort_processing, pattern='^abort$'))
        application.add_handler(CallbackQueryHandler(button_callback, pattern='^(browse|file):'))

        logging.info("Bot is running...")

        asyncio.get_event_loop().create_task(self.check_doreboot_file(application))

        application.run_polling()