Update bot.py

bot.py

@@ -2,6 +2,7 @@
 
 import logging
 import os
+import traceback
 from telegram import Update, ReplyKeyboardMarkup, ReplyKeyboardRemove, InlineKeyboardButton, InlineKeyboardMarkup
 from telegram.ext import (
     Application,
@@ -22,122 +23,187 @@ logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(level
 logger = logging.getLogger(__name__)
 
 # --- State Definitions for Conversation ---
-# We define states to keep track of where the user is in the conversation.
-SELECTING_ACTION, SELECTING_STYLE, SELECTING_GENDER = range(3)
+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 ---
-# This loads the heavy AI models into memory once when the bot starts.
 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}")
-    # In a real deployment, you might want to exit or send an alert here.
     generation_service = None
 
 
-# --- Command Handlers ---
+# --- 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)
 
-async def start_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
-    """
-    Handles the /start command. Greets the user, creates their account in the DB if needed,
-    and shows them their credit balance.
-    """
+# --- 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.")
-    
-    # Get or create user in the database
-    db_user = database.get_or_create_user(user.id, user.username)
-    
-    if not db_user:
-        await update.message.reply_text("😕 Sorry, there was a problem setting up your account. Please try again later.")
-        return ConversationHandler.END
-
-    credits = db_user.get('credits', 0)
+    database.get_or_create_user(user.id, user.username)
+    credits = database.get_user_credits(user.id)
     
     welcome_message = (
-        f"👋 Welcome to the Magic Face Bot, {user.first_name}!\n\n"
-        "I can transform your face into legendary characters using AI.\n\n"
-        f"You currently have **{credits}** credits.\n\n"
-        "To get started, just send me a clear photo of your face."
+        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")
     
-    # We don't enter a state here yet; we wait for the user to send a photo.
-    return SELECTING_ACTION
+    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 credits_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
-    """
-    Handles the /credits command. Shows the user their current credit balance.
-    """
+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
-    credits = database.get_user_credits(user_id)
-    await update.message.reply_text(f"You have **{credits}** credits remaining.", parse_mode='Markdown')
-
-
-async def help_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
-    """
-    Handles the /help command.
-    """
-    help_text = (
-        "**How to use the Magic Face Bot:**\n\n"
-        "1. Send a clear, well-lit photo of your face.\n"
-        "2. Choose a transformation style from the menu.\n"
-        "3. Select your gender.\n"
-        "4. Wait for the AI to work its magic!\n\n"
-        "**Commands:**\n"
-        "/start - Restart the bot.\n"
-        "/credits - Check your credit balance.\n"
-        "/help - Show this message."
-    )
-    await update.message.reply_text(help_text, parse_mode='Markdown')
+    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:
-    """
-    Starts the bot.
-    """
     if not generation_service:
         logger.error("Generation service not available. Bot is shutting down.")
         return
 
-    # Create the Application and pass it your bot's token.
     application = Application.builder().token(config.TELEGRAM_TOKEN).build()
 
-    # --- Conversation Handler (to be built out) ---
-    # For now, we just have the entry point and basic commands.
     conv_handler = ConversationHandler(
-        entry_points=[
-            CommandHandler("start", start_command),
-            # We will add a MessageHandler for photos here later
-        ],
+        entry_points=[MessageHandler(filters.PHOTO, photo_handler)],
         states={
-            SELECTING_ACTION: [
-                # Handlers for when the user sends a photo will go here
-            ],
-            SELECTING_STYLE: [
-                # Handlers for when the user chooses a style
-            ],
-            SELECTING_GENDER: [
-                # Handlers for when the user chooses a gender
-            ],
+            SELECTING_STYLE: [CallbackQueryHandler(style_selected_handler, pattern="^style_")],
+            SELECTING_GENDER: [CallbackQueryHandler(gender_selected_handler, pattern="^gender_")],
         },
-        fallbacks=[CommandHandler("start", start_command)],
+        fallbacks=[CommandHandler("start", start_command), CommandHandler("cancel", cancel_command)],
     )
 
+    application.add_handler(CommandHandler("start", start_command))
     application.add_handler(conv_handler)
-    application.add_handler(CommandHandler("help", help_command))
-    application.add_handler(CommandHandler("credits", credits_command))
+    # Other handlers can be added here
 
-    # Run the bot until the user presses Ctrl-C
     logger.info("Bot is starting to poll for updates...")
     application.run_polling()
 
 
 if __name__ == "__main__":
     main()
-