# bot.py import logging import os import traceback from telegram import Update, ReplyKeyboardMarkup, ReplyKeyboardRemove, InlineKeyboardButton, InlineKeyboardMarkup from telegram.ext import ( Application, CommandHandler, ContextTypes, ConversationHandler, MessageHandler, CallbackQueryHandler, filters, ) import config import database from generate import GenerationService # --- Setup Logging --- logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s") logger = logging.getLogger(__name__) # --- State Definitions for Conversation --- SELECTING_STYLE, SELECTING_GENDER = range(2) # --- Style Presets --- # These are the options the user will see. STYLE_PRESETS = [ {"title": "Mona Lisa", "prompt": "A mesmerizing portrait in the style of Leonardo da Vinci's Mona Lisa, renaissance oil painting, soft sfumato technique", "emoji": "🎨"}, {"title": "Iron Hero", "prompt": "Hyper realistic portrait as a high-tech superhero, wearing advanced metallic suit, arc reactor glow", "emoji": "🦾"}, {"title": "Van Gogh", "prompt": "Self-portrait in the style of Vincent van Gogh, bold brushstrokes, vibrant colors, post-impressionist style", "emoji": "🎨"}, {"title": "Jedi Master", "prompt": "Epic portrait as a Jedi Master from Star Wars, wearing traditional robes, holding a lightsaber", "emoji": "⚔️"}, {"title": "Greek God", "prompt": "Mythological portrait as a Greek God, wearing flowing robes, golden laurel wreath, Mount Olympus background", "emoji": "⚡"}, {"title": "Medieval Knight", "prompt": "Noble knight portrait, wearing ornate plate armor, holding a sword, castle background", "emoji": "🛡️"}, ] # --- Initialize the Generation Service --- logger.info("Initializing Generation Service... This may take a while.") try: generation_service = GenerationService() logger.info("Generation Service initialized successfully.") except Exception as e: logger.error(f"FATAL: Could not initialize GenerationService. Bot cannot start. Error: {e}") generation_service = None # --- Helper Functions --- def get_style_keyboard() -> InlineKeyboardMarkup: """Creates the inline keyboard for style selection.""" keyboard = [ [InlineKeyboardButton(f"{style['emoji']} {style['title']}", callback_data=f"style_{i}")] for i, style in enumerate(STYLE_PRESETS) ] return InlineKeyboardMarkup(keyboard) # --- Conversation Handlers --- async def start_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None: """Handles the /start command.""" user = update.effective_user logger.info(f"User {user.id} ({user.username}) started the bot.") database.get_or_create_user(user.id, user.username) credits = database.get_user_credits(user.id) welcome_message = ( f"👋 Welcome, {user.first_name}!\n\n" "To transform your face, just send me a clear photo.\n\n" f"You have **{credits}** credits." ) await update.message.reply_text(welcome_message, parse_mode='Markdown') async def photo_handler(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int: """Handles the user sending a photo.""" user = update.effective_user photo_file = await update.message.photo[-1].get_file() # Create a temporary directory to store the photo temp_dir = "temp_user_photos" os.makedirs(temp_dir, exist_ok=True) file_path = os.path.join(temp_dir, f"{user.id}_{photo_file.file_unique_id}.jpg") await photo_file.download_to_drive(file_path) logger.info(f"Photo from user {user.id} saved to {file_path}") context.user_data['photo_path'] = file_path await update.message.reply_text( "Great photo! Now, pick a style for your transformation:", reply_markup=get_style_keyboard() ) return SELECTING_STYLE async def style_selected_handler(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int: """Handles the user selecting a style from the inline keyboard.""" query = update.callback_query await query.answer() # Acknowledge the button press style_index = int(query.data.split('_')[1]) selected_style = STYLE_PRESETS[style_index] context.user_data['prompt'] = selected_style['prompt'] logger.info(f"User {update.effective_user.id} selected style: {selected_style['title']}") keyboard = [ [ InlineKeyboardButton("♀️ Female", callback_data="gender_Female"), InlineKeyboardButton("♂️ Male", callback_data="gender_Male"), ] ] await query.edit_message_text( text=f"Style selected: *{selected_style['title']}*.\n\nNow, please select your gender:", reply_markup=InlineKeyboardMarkup(keyboard), parse_mode='Markdown' ) return SELECTING_GENDER async def gender_selected_handler(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int: """Handles the user selecting a gender and triggers the final generation.""" query = update.callback_query await query.answer() user_id = update.effective_user.id gender = query.data.split('_')[1] context.user_data['gender'] = gender logger.info(f"User {user_id} selected gender: {gender}. Preparing for generation.") # Check for credits before starting if database.get_user_credits(user_id) <= 0: await query.edit_message_text("Oh no! You're out of credits. Use /buy to get more.") return ConversationHandler.END await query.edit_message_text("✨ Magic in progress... Your masterpiece is being created. This can take several minutes!") try: # Get all data from context photo_path = context.user_data['photo_path'] prompt = context.user_data['prompt'] # Call the generation service image_url = generation_service.generate_magic_image( face_images=[photo_path], gender=gender, prompt=prompt ) if image_url: logger.info(f"Successfully generated image for user {user_id}. URL: {image_url}") database.deduct_credit(user_id) new_credits = database.get_user_credits(user_id) await context.bot.send_photo( chat_id=user_id, photo=image_url, caption=f"Here is your masterpiece! ✨\n\nYou have {new_credits} credits left." ) else: logger.error(f"Image generation failed for user {user_id}.") await context.bot.send_message(user_id, "😕 Something went wrong during the magic spell. Please try again with a different photo.") except Exception as e: logger.error(f"An error occurred in the generation process for user {user_id}: {e}") traceback.print_exc() await context.bot.send_message(user_id, "Sorry, a critical error occurred. The developer has been notified.") finally: # Clean up the temporary photo if 'photo_path' in context.user_data and os.path.exists(context.user_data['photo_path']): os.remove(context.user_data['photo_path']) # End the conversation return ConversationHandler.END async def cancel_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int: """Cancels and ends the conversation.""" await update.message.reply_text("Action cancelled.", reply_markup=ReplyKeyboardRemove()) return ConversationHandler.END # --- Main Application --- def main() -> None: if not generation_service: logger.error("Generation service not available. Bot is shutting down.") return application = Application.builder().token(config.TELEGRAM_TOKEN).build() conv_handler = ConversationHandler( entry_points=[MessageHandler(filters.PHOTO, photo_handler)], states={ SELECTING_STYLE: [CallbackQueryHandler(style_selected_handler, pattern="^style_")], SELECTING_GENDER: [CallbackQueryHandler(gender_selected_handler, pattern="^gender_")], }, fallbacks=[CommandHandler("start", start_command), CommandHandler("cancel", cancel_command)], ) application.add_handler(CommandHandler("start", start_command)) application.add_handler(conv_handler) # Other handlers can be added here logger.info("Bot is starting to poll for updates...") application.run_polling() if __name__ == "__main__": main()