Create bot.py

bot.py

@@ -0,0 +1,143 @@
+# bot.py
+
+import logging
+import os
+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 ---
+# We define states to keep track of where the user is in the conversation.
+SELECTING_ACTION, SELECTING_STYLE, SELECTING_GENDER = range(3)
+
+# --- 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 ---
+
+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.
+    """
+    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)
+    
+    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."
+    )
+    
+    await update.message.reply_text(welcome_message, parse_mode='Markdown')
+    
+    # We don't enter a state here yet; we wait for the user to send a photo.
+    return SELECTING_ACTION
+
+
+async def credits_command(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
+    """
+    Handles the /credits command. Shows the user their current credit balance.
+    """
+    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')
+
+
+# --- 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
+        ],
+        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
+            ],
+        },
+        fallbacks=[CommandHandler("start", start_command)],
+    )
+
+    application.add_handler(conv_handler)
+    application.add_handler(CommandHandler("help", help_command))
+    application.add_handler(CommandHandler("credits", credits_command))
+
+    # 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()
+