import os
from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes, CallbackQueryHandler

class TelegramHelper:
    def __init__(self, bot):
        self.bot = bot
        self.TELEGRAM_BOT_TOKEN = os.getenv('TELEGRAM_BOT_TOKEN')

    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? You can send me images and then ask questions about them."
        )

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

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

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

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

        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)
        await update.message.reply_text(response)

    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
        abort_message = await self.bot.abort_processing(user_id)
        await query.edit_message_text(text=abort_message)

    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
        )

    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("status", self.status))
        application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, self.handle_message))
        application.add_handler(CallbackQueryHandler(self.abort_processing, pattern='^abort$'))

        application.run_polling()