bot.py

# -----------------------------------------------------------------------------
# Copyright (c) 2025 tanbaycu. All rights reserved.
# Licensed under the MIT License. See LICENSE file for details.
# -----------------------------------------------------------------------------

import logging
import asyncio
import json
import requests
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

from telegram import Update, ReplyKeyboardMarkup, ReplyKeyboardRemove
from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes, ConversationHandler
from form_parser import get_form_data
from ai_generator import setup_gemini, generate_responses
import database
from payments.webhook_server import start_server
from datetime import time, timezone, timedelta


# Enable logging
logging.basicConfig(
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO
)
logger = logging.getLogger(__name__)

# States
LINK, COUNT = range(2)

# Bot Token
TOKEN = os.getenv("TELEGRAM_TOKEN")
if not TOKEN:
    logger.error("TELEGRAM_TOKEN not found in .env file!")
    raise ValueError("TELEGRAM_TOKEN not found in .env file!")

# Shared Dictionary to store user data in memory
# Key: UserID, Value: {'url': ..., 'form_data': ...}
user_context_store = {}

async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    """Sends a welcome message."""
    user = update.effective_user
    await update.message.reply_text(
        f"👋 *Xin chào {user.first_name}!* \n\n"
        "Tôi là **Auto Form Bot Pro**. 🤖\n"
        "Tôi giúp bạn điền Google Form tự động với tốc độ bàn thờ.\n\n"
        "📜 *Danh sách lệnh:*\n"
        "• /submit - Bắt đầu điền Form (hoặc `/submit <link>`)\n"
        "• /config - Cài đặt tốc độ & AI\n"
        "• /persona - Cài đặt danh tính\n"
        "• /balance - Xem số dư Xu\n"
        "• /deposit - Nạp tiền tự động\n"
        "• /price - Bảng giá dịch vụ\n"
        "• /help - Hướng dẫn chi tiết\n\n"
        "👇 *Bấm vào Menu hoặc gõ /submit để bắt đầu!*",
        parse_mode="Markdown"
    )
    # Ensure user exists in DB
    database.get_user(user.id, user.first_name)
    return ConversationHandler.END

from utils import safe_delete, get_progress_bar

async def submit_entry(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    """Entry point for /submit command."""
    # Check if user provided link arguments: /submit <url>
    if context.args:
        url = context.args[0]
        # Inject URL into message text so analyze_link can find it
        update.message.text = url
        return await analyze_link(update, context)
    
    msg = await update.message.reply_text(
        "🚀 *Bắt đầu Task mới*\n\n"
        "Vui lòng gửi **Link Google Form** cho tôi.\n"
        "_(Hỗ trợ form nhiều trang)_",
        parse_mode="Markdown"
    )
    # Store message ID to delete later if needed (optional UX)
    context.user_data['last_bot_msg'] = msg.message_id
    return LINK

async def analyze_link(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    """Analyzes the provided link."""
    url = update.message.text
    user_id = update.effective_user.id
    
    # Cleaning up previous prompt if exists
    if 'last_bot_msg' in context.user_data:
        await safe_delete(context, update.effective_chat.id, context.user_data['last_bot_msg'])
    
    msg_processing = await update.message.reply_text("🔍 *Đang phân tích cấu trúc Form...*", parse_mode="Markdown")
    
    loop = asyncio.get_running_loop()
    form_data = await loop.run_in_executor(None, get_form_data, url)
    
    if not form_data or not form_data['fields']:
        await context.bot.edit_message_text(
            chat_id=update.effective_chat.id,
            message_id=msg_processing.message_id,
            text="❌ *Lỗi phân tích:*\n"
                 "- Link không đúng định dạng?\n"
                 "- Form yêu cầu đăng nhập?\n"
                 "🔄 Gõ /submit để thử lại.",
            parse_mode="Markdown"
        )
        return ConversationHandler.END
    
    # Store data
    user_context_store[user_id] = {
        'url': url,
        'form_data': form_data
    }
    
    # Summarize fields
    fields = form_data['fields']
    field_counts = {}
    for f in fields:
        f_type = f.get('type', 'unknown')
        field_counts[f_type] = field_counts.get(f_type, 0) + 1
    
    type_summary = ", ".join([f"{k}: {v}" for k, v in field_counts.items()])

    msg = f"✅ *Phân tích thành công!*\n\n"
    msg += f"📄 **Tiêu đề:** `{form_data['title']}`\n"
    msg += f"📊 **Câu hỏi:** {len(fields)} ({type_summary})\n\n"
    msg += "🔢 *Nhập số lượng muốn điền:* (Ví dụ: `10`)"
    
    await context.bot.edit_message_text(
        chat_id=update.effective_chat.id,
        message_id=msg_processing.message_id,
        text=msg,
        parse_mode="Markdown"
    )
    context.user_data['last_bot_msg'] = msg_processing.message_id
    return COUNT

async def perform_filling(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    """Generates content and submits the form."""
    user_id = update.effective_user.id
    text = update.message.text
    
    if not text.isdigit():
        await update.message.reply_text("❌ Vui lòng nhập một số hợp lệ (ví dụ: 10, 20).")
        return COUNT
    
    count = int(text)
    user_data = user_context_store.get(user_id)
    
    if not user_data:
        await update.message.reply_text("⚠️ Có lỗi dữ liệu, vui lòng bắt đầu lại với /start")
        return ConversationHandler.END
    
    if count > 50:
         await update.message.reply_text("⚠️ Số lượng quá lớn, tôi chỉ hỗ trợ tối đa **50** lần một lúc để an toàn.", parse_mode="Markdown")
         count = 50

    # --- COST CALCULATION (DISABLED) ---
    # cost_per_unit = 0.1
    # total_cost = count * cost_per_unit
    
    # db_user = database.get_user(user_id, update.effective_user.first_name)
    # current_balance = db_user['balance']
    
    # if current_balance < total_cost:
    #     missing = total_cost - current_balance
    #     await update.message.reply_text(
    #         f"🚫 **Số dư không đủ!**\n\n"
    #         f"💰 Số dư hiện tại: **{current_balance:.1f} Xu**\n"
    #         f"📉 Chi phí yêu cầu: **{total_cost:.1f} Xu**\n"
    #         f"⚠️ Còn thiếu: **{missing:.1f} Xu**\n\n"
    #         "👉 Vui lòng nạp thêm bằng lệnh /deposit",
    #         parse_mode="Markdown"
    #     )
    #     return ConversationHandler.END

    # Deduct Balance
    # database.update_balance(user_id, -total_cost, f"Điền Form: {user_data['form_data']['title'][:20]}...")
    
    await update.message.reply_text(f"🚀 *Bắt đầu nhiệm vụ:* Điền {count} lần.\n⏳ _Vui lòng không spam lệnh khác trong lúc chờ..._", parse_mode="Markdown")
    
    model = setup_gemini()
    action_url = user_data['form_data']['action_url']
    fields = user_data['form_data']['fields']
    
    success_count = 0
    
    status_message = await update.message.reply_text(
        f"⏳ *Tiến độ:* 0/{count}\n" + get_progress_bar(0, count) + "\n\n🤖 Đang khởi động AI...",
        parse_mode="Markdown"
    )
    
    # Initialize Stop Signal
    user_context_store[user_id]['stop_signal'] = False
    
    from commands.config_cmd import get_user_config, CONFIG_DELAY
    user_cfg = get_user_config(user_id)
    delay_sec = user_cfg.get(CONFIG_DELAY, 2)

    for i in range(count):
        # CHECK STOP SIGNAL
        if user_context_store.get(user_id, {}).get('stop_signal'):
            await update.message.reply_text("🛑 Đã dừng theo yêu cầu của bạn.")
            break

        current_status_text = ""
        try:
            # 1. Get Persona & Imported Data
            from commands.persona_cmd import get_user_persona
            from commands.import_cmd import get_data_row
            
            user_persona = get_user_persona(user_id)
            imported_row = get_data_row(user_id, i) # i is the current iteration index

            # 2. Generate Answers
            answers = generate_responses(user_data['form_data'], model, persona=user_persona, imported_data=imported_row)
            
            # 3. Prepare Payload
            payload = {}
            # Add hidden fields (fvv, etc.) but exclude potentially problematic ones
            if 'hidden_fields' in user_data['form_data']:
                for k, v in user_data['form_data']['hidden_fields'].items():
                     if k not in ['partialResponse', 'submissionTimestamp']:
                         payload[k] = v

            for field in fields:
                f_id = field['id']
                # Google forms use entry.ID as key
                entry_key = f"entry.{f_id}"
                
                # Check if we generated an answer
                if f_id in answers:
                    val = answers[f_id]
                    # Handle lists (checkboxes) by assigning directly
                    payload[entry_key] = val
            
            # 3. Submit
            headers = {
                 "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
            }
            
            # Proxy Selection
            from commands.proxy_cmd import get_random_proxy
            proxy_dict = get_random_proxy()
            if proxy_dict:
                logger.info(f"Using Proxy: {proxy_dict['http']}")
            
            logger.info(f"Submitting payload with headers")
            
            # Use proxied request NON-BLOCKING
            def send_request():
                return requests.post(action_url, data=payload, headers=headers, proxies=proxy_dict, timeout=20)
            
            loop = asyncio.get_running_loop()
            submit_resp = await loop.run_in_executor(None, send_request)
            
            if submit_resp.status_code == 200:
                success_count += 1
                current_status_text = "✅ Đã gửi"
            else:
                logger.error(f"Form submission failed: {submit_resp.status_code}")
                # Log a snippet of response
                logger.error(f"Response: {submit_resp.text[:500]}")
                current_status_text = f"❌ Lỗi {submit_resp.status_code}"
                
        except Exception as e:
            logger.error(f"Error during loop: {e}")
            current_status_text = "❌ Lỗi hệ thống"
        
        # Update progress every few items or at specific milestones
        update_interval = 1 if count <= 5 else 5
        
        if (i + 1) % update_interval == 0 or i == count - 1:
             progress_bar = get_progress_bar(i + 1, count)
             try:
                await context.bot.edit_message_text(
                    chat_id=update.effective_chat.id,
                    message_id=status_message.message_id,
                    text=f"⏳ *Tiến độ:* {i + 1}/{count}\n{progress_bar}\n\n🤖 Tình trạng: {current_status_text}",
                    parse_mode="Markdown"
                )
             except Exception as e:
                 logger.warning(f"Could not update status message: {e}")
        
        # DELAY to avoid spam
        await asyncio.sleep(delay_sec)

    report_msg = f"🎉 *Nhiệm vụ hoàn tất!*\n\n"
    report_msg += f"✅ **Thành công:** {success_count}/{count}\n"
    if success_count < count:
        report_msg += f"❌ **Thất bại:** {count - success_count}\n"
    report_msg += "\nCảm ơn bạn đã sử dụng dịch vụ! Gõ /start để điền form khác."

    # Save to history
    from commands.history_cmd import add_history
    add_history(user_id, user_data['form_data']['title'], success_count, count)

    await update.message.reply_text(report_msg, parse_mode="Markdown")
    return ConversationHandler.END

async def cancel(update: Update, context: ContextTypes.DEFAULT_TYPE) -> int:
    """Cancels and ends the conversation."""
    await update.message.reply_text(
        "🚫 Đã hủy thao tác. Gõ /start để bắt đầu lại.", reply_markup=ReplyKeyboardRemove()
    )
    return ConversationHandler.END

async def daily_reminder(context: ContextTypes.DEFAULT_TYPE):
    """Sends a daily reminder to all users."""
    logger.info("Running daily reminder job...")
    try:
        user_ids = database.get_all_users()
        for user_id in user_ids:
            try:
                await context.bot.send_message(
                    chat_id=user_id,
                    text="🌞 *Chào buổi sáng!* \n\nBạn có Form nào cần điền hôm nay không?\n👉 Gõ /submit để bắt đầu ngay nhé!",
                    parse_mode="Markdown"
                )
                # Sleep briefly to avoid hitting broadcast limits
                await asyncio.sleep(0.1)
            except Exception as e:
                # User might have blocked the bot
                logger.warning(f"Failed to send reminder to {user_id}: {e}")
    except Exception as e:
        logger.error(f"Error in daily_reminder: {e}")

def main() -> None:
    """Run the bot."""
    from input_handlers import global_cancel_command
    # Build application
    application = Application.builder().token(TOKEN).build()

    # Register Command Handlers
    

    # 4. Config
    from commands.config_cmd import get_config_handlers
    for handler in get_config_handlers():
        application.add_handler(handler)
        
    # 5. Stats
    from commands.stats_cmd import get_stats_handlers
    for handler in get_stats_handlers():
        application.add_handler(handler)
        
    # 6. Stop
    from commands.stop_cmd import get_stop_handlers
    for handler in get_stop_handlers():
        application.add_handler(handler)

    # 7. Proxy
    from commands.proxy_cmd import get_proxy_handlers
    for handler in get_proxy_handlers():
        application.add_handler(handler)

    # 8. Import
    from commands.import_cmd import get_import_handlers
    for handler in get_import_handlers():
        application.add_handler(handler)

    # 9. Persona
    from commands.persona_cmd import get_persona_handlers
    for handler in get_persona_handlers():
        application.add_handler(handler)

    # 10. History
    from commands.history_cmd import get_history_handlers
    for handler in get_history_handlers():
        application.add_handler(handler)

    # 11. Help
    from commands.help_cmd import get_help_handlers
    for handler in get_help_handlers():
        application.add_handler(handler)
        
    # 12. Payment
    from commands.payment_cmd import get_payment_handlers
    for handler in get_payment_handlers():
        application.add_handler(handler)

    # Conversation handler (Main Flow)
    conv_handler = ConversationHandler(
        entry_points=[
            CommandHandler("submit", submit_entry),
        ],
        states={
            LINK: [MessageHandler(filters.TEXT & ~filters.COMMAND, analyze_link)],
            COUNT: [MessageHandler(filters.TEXT & ~filters.COMMAND, perform_filling)],
        },
        fallbacks=[
            CommandHandler("cancel", cancel),
            MessageHandler(filters.COMMAND, global_cancel_command)
        ],
    )

    application.add_handler(conv_handler)
    application.add_handler(CommandHandler("start", start))
    # AUTO-UPDATE BOT COMMANDS MENU
    from telegram import BotCommand
    commands = [
        BotCommand("submit", "🚀 Điền form mới"),
        BotCommand("start", "👋 Giới thiệu Bot"),
        BotCommand("stop", "🛑 Dừng khẩn cấp"),
        BotCommand("config", "⚙️ Cài đặt (Delay/AI)"),
        BotCommand("proxy", "🛡️ Quản lý Proxy"),
        BotCommand("import", "📂 Nạp dữ liệu (CSV/Excel)"),
        BotCommand("persona", "👤 Cấu hình danh tính"),
        BotCommand("stats", "📊 Thống kê"),
        BotCommand("history", "📜 Lịch sử"),
        BotCommand("history", "📜 Lịch sử"),
        BotCommand("balance", "💰 Xem số dư (Dev)"),
        BotCommand("deposit", "💳 Nạp tiền (Dev)"),
        BotCommand("price", "💲 Bảng giá (Dev)"),
        BotCommand("help", "📘 Hướng dẫn"),
    ]
    
    async def post_init(app: Application):
        await app.bot.set_my_commands(commands)
        
        # Start Webhook Server in background
        asyncio.create_task(start_server())
        
        # Schedule Daily Reminder at 8:00 AM (UTC+7)
        t = time(8, 0, tzinfo=timezone(timedelta(hours=7)))
        app.job_queue.run_daily(daily_reminder, t)
        
        print(f"Daily reminder scheduled for {t}")
        
        print("Bot Commands Menu Updated & Webhook Started!")

    print("Bot is running...")
    application.post_init = post_init
    application.run_polling(allowed_updates=Update.ALL_TYPES)

if __name__ == "__main__":
    database.init_db()
    main()