telegram_bot.py

from telegram.ext import Application, CommandHandler, CallbackQueryHandler, ContextTypes, ConversationHandler, MessageHandler, filters
from telegram import InlineKeyboardButton, InlineKeyboardMarkup, Update
import asyncio
import threading
from main import PumpBot
import json
import random
from token_generator import CommentGenerator
import time
import requests
from dotenv import load_dotenv
import os

# Load environment variables
load_dotenv()

# Define conversation states
CHOOSING, INPUT_CHANNEL, GENERATE_COMMENTS = range(3)  # Removed unused states

class BotTask:
    def __init__(self, token_address, username, wallet, status="running"):
        self.token_address = token_address
        self.username = username
        self.wallet = wallet
        self.status = status  # running, paused, stopped
        self.created_at = time.time()
        self.comments_made = 0

class UserPool:
    def __init__(self, max_pool_size=50, min_pool_size=10):
        self.users = []  # List of user objects
        self.max_pool_size = max_pool_size
        self.min_pool_size = min_pool_size
        self.max_comments_per_user = 50  # Max lifetime comments
        self.cooldown_period = 3600  # 1 hour between comments
        self.min_cooldown = 1800  # 30 min absolute minimum
        self.load_users()

    def get_available_user(self, force_new=False):
        """Get an available user based on smart selection"""
        current_time = time.time()
        available_users = []

        # First try to find ideal users (good cooldown, low usage)
        for user in self.users:
            if (user['status'] == 'available' and 
                current_time - user['last_used'] > self.cooldown_period and 
                user['comments_made'] < self.max_comments_per_user):
                available_users.append(user)

        # If we have ideal users, pick randomly from them
        if available_users and not force_new:
            return random.choice(available_users)

        # If pool is below min size or force_new, create new user
        if len(self.users) < self.min_pool_size or force_new:
            return None

        # Otherwise, find least recently used user with minimum cooldown
        least_recent = min(
            [u for u in self.users if u['status'] == 'available'],
            key=lambda x: x['last_used'],
            default=None
        )
        
        if least_recent and current_time - least_recent['last_used'] > self.min_cooldown:
            return least_recent

        return None

    def add_user(self, wallet, username, bio):
        """Add new user to pool with health tracking"""
        user = {
            'wallet': wallet,
            'username': username,
            'bio': bio,
            'created_at': time.time(),
            'last_used': time.time(),
            'comments_made': 0,
            'status': 'available',
            'health_score': 100,  # Track account health
            'success_rate': 100,  # Track successful comments
            'errors': []  # Track error history
        }
        self.users.append(user)
        self.save_users()

        # Remove oldest unhealthy users if pool too large
        if len(self.users) > self.max_pool_size:
            self.cleanup_pool()

    def release_user(self, user, success=True, error=None):
        """Update user stats after usage"""
        user['status'] = 'available'
        user['last_used'] = time.time()
        
        if success:
            user['comments_made'] += 1
            user['success_rate'] = ((user['success_rate'] * (user['comments_made'] - 1) + 100) / 
                                  user['comments_made'])
        else:
            user['success_rate'] = ((user['success_rate'] * user['comments_made']) / 
                                  (user['comments_made'] + 1))
            if error:
                user['errors'].append({
                    'time': time.time(),
                    'error': str(error)
                })

        # Update health score based on various factors
        age_days = (time.time() - user['created_at']) / 86400
        user['health_score'] = min(100, (
            user['success_rate'] * 0.4 +  # 40% weight on success rate
            (1 - (user['comments_made'] / self.max_comments_per_user)) * 30 +  # 30% weight on usage
            min(age_days / 7, 1) * 30  # 30% weight on account age (max 7 days)
        ))

        self.save_users()

    def cleanup_pool(self):
        """Remove unhealthy or overused accounts"""
        # Sort by health score and remove worst ones
        self.users.sort(key=lambda x: x['health_score'])
        while len(self.users) > self.max_pool_size:
            removed = self.users.pop(0)
            print(f"Removed unhealthy user {removed['username']} (health: {removed['health_score']:.1f})")
        self.save_users()

    def get_pool_stats(self):
        """Get statistics about the user pool"""
        return {
            'total_users': len(self.users),
            'available_users': len([u for u in self.users if u['status'] == 'available']),
            'healthy_users': len([u for u in self.users if u['health_score'] >= 70]),
            'avg_health': sum(u['health_score'] for u in self.users) / len(self.users) if self.users else 0,
            'avg_success_rate': sum(u['success_rate'] for u in self.users) / len(self.users) if self.users else 0,
            'total_comments': sum(u['comments_made'] for u in self.users)
        }

    def save_users(self):
        """Save user pool to file"""
        with open('user_pool.json', 'w') as f:
            json.dump(self.users, f)
            
    def load_users(self):
        """Load user pool from file"""
        try:
            with open('user_pool.json', 'r') as f:
                self.users = json.load(f)
        except:
            self.users = []

class TelegramManager:
    def __init__(self):
        self.telegram_token = os.getenv('TELEGRAM_BOT_TOKEN')
        if not self.telegram_token:
            raise ValueError("TELEGRAM_BOT_TOKEN not found in environment variables")
            
        self.active_bots = {}  # Store running bots by user_id
        self.captcha_api_key = os.getenv('CAPTCHA_API_KEY')
        if not self.captcha_api_key:
            raise ValueError("CAPTCHA_API_KEY not found in environment variables")
            
        self.current_token = None  # Store token temporarily for comment generation
        
        openai_api_key = os.getenv('OPENAI_API_KEY')
        if not openai_api_key:
            raise ValueError("OPENAI_API_KEY not found in environment variables")
            
        self.comment_generator = CommentGenerator(openai_api_key)
        self.tasks = {}  # user_id -> list of BotTask objects
        self.user_pool = UserPool(max_pool_size=50)
        
    def get_main_keyboard(self):
        """Get the main menu keyboard"""
        keyboard = [
            [InlineKeyboardButton("💭 Generate Bullish Comments", callback_data='generate_comments')],
            [InlineKeyboardButton("📢 Promote Channel", callback_data='promote_channel')],
            [InlineKeyboardButton("📋 View Tasks", callback_data='view_tasks')],
            [InlineKeyboardButton("❌ Stop All Bots", callback_data='stop_bots')]
        ]
        return InlineKeyboardMarkup(keyboard)

    def get_back_keyboard(self):
        """Get keyboard with back button"""
        keyboard = [[InlineKeyboardButton("⬅️ Back to Menu", callback_data='back_to_menu')]]
        return InlineKeyboardMarkup(keyboard)

    def get_comment_quantity_keyboard(self):
        """Get keyboard for comment quantity selection"""
        keyboard = [
            [InlineKeyboardButton("5 Comments 💬", callback_data='gen_5')],
            [InlineKeyboardButton("10 Comments 💬", callback_data='gen_10')],
            [InlineKeyboardButton("15 Comments 💬", callback_data='gen_15')],
            [InlineKeyboardButton("⬅️ Back to Menu", callback_data='back_to_menu')]
        ]
        return InlineKeyboardMarkup(keyboard)

    def get_task_management_keyboard(self, task_id):
        """Get keyboard for managing a specific task"""
        keyboard = [
            [
                InlineKeyboardButton("⏸️ Pause", callback_data=f'pause_task_{task_id}'),
                InlineKeyboardButton("▶️ Resume", callback_data=f'resume_task_{task_id}'),
                InlineKeyboardButton("⏹️ Stop", callback_data=f'stop_task_{task_id}')
            ],
            [InlineKeyboardButton("⬅️ Back to Tasks", callback_data='view_tasks')],
            [InlineKeyboardButton("🏠 Main Menu", callback_data='back_to_menu')]
        ]
        return InlineKeyboardMarkup(keyboard)

    async def start(self, update, context):
        """Send welcome message with main menu"""
        await update.message.reply_text(
            "👋 Welcome to the Pump Bot Manager!\n\n"
            "What would you like to do?",
            reply_markup=self.get_main_keyboard()
        )
        return CHOOSING

    async def button_handler(self, update, context):
        """Handle button presses"""
        query = update.callback_query
        await query.answer()

        if query.data == 'back_to_menu':
            await query.message.edit_text(
                "What would you like to do?",
                reply_markup=self.get_main_keyboard()
            )
            return CHOOSING

        elif query.data == 'promote_channel':
            await query.message.edit_text(
                "📝 Please send the Telegram channel you want to promote:\n"
                "Example: @channel_name",
                reply_markup=self.get_back_keyboard()
            )
            return INPUT_CHANNEL

        elif query.data == 'generate_comments':
            await query.message.edit_text(
                "📝 Please send the token address you want to generate comments for:",
                reply_markup=self.get_back_keyboard()
            )
            return GENERATE_COMMENTS

        elif query.data.startswith('gen_'):
            quantity = int(query.data.split('_')[1])
            if self.current_token:
                await query.message.edit_text(
                    f"🤖 Starting {quantity} bots to comment on {self.current_token}...\n"
                    "This may take a few minutes."
                )
                
                # Create and start bots in a separate thread to not block
                thread = threading.Thread(
                    target=self.run_multiple_bots,
                    args=(query.message, self.current_token, quantity)
                )
                thread.start()
                
                # Reset current token
                self.current_token = None
                return CHOOSING

        elif query.data == 'view_tasks':
            user_id = query.from_user.id
            if user_id not in self.tasks or not self.tasks[user_id]:
                await query.message.edit_text(
                    "No active tasks found!\n\n"
                    "What would you like to do?",
                    reply_markup=self.get_main_keyboard()
                )
                return CHOOSING

            # Show list of tasks
            task_list = "📋 Your Active Tasks:\n\n"
            for i, task in enumerate(self.tasks[user_id], 1):
                runtime = time.time() - task.created_at
                hours = int(runtime // 3600)
                minutes = int((runtime % 3600) // 60)
                
                task_list += (
                    f"{i}. Token: {task.token_address}\n"
                    f"   Bot: {task.username}\n"
                    f"   Status: {task.status}\n"
                    f"   Runtime: {hours}h {minutes}m\n"
                    f"   Comments: {task.comments_made}\n\n"
                )

            keyboard = [
                [InlineKeyboardButton(f"Manage Task {i}", callback_data=f'manage_task_{i-1}')]
                for i in range(1, len(self.tasks[user_id]) + 1)
            ]
            keyboard.append([InlineKeyboardButton("🏠 Main Menu", callback_data='back_to_menu')])
            
            await query.message.edit_text(
                task_list,
                reply_markup=InlineKeyboardMarkup(keyboard)
            )
            return CHOOSING

        elif query.data.startswith('manage_task_'):
            task_id = int(query.data.split('_')[-1])
            user_id = query.from_user.id
            task = self.tasks[user_id][task_id]
            
            task_info = (
                f"📊 Task Details:\n\n"
                f"Token: {task.token_address}\n"
                f"Bot: {task.username}\n"
                f"Status: {task.status}\n"
                f"Comments Made: {task.comments_made}\n"
            )
            
            await query.message.edit_text(
                task_info,
                reply_markup=self.get_task_management_keyboard(task_id)
            )
            return CHOOSING

        elif query.data.startswith(('pause_task_', 'resume_task_', 'stop_task_')):
            action, task_id = query.data.split('_task_')
            task_id = int(task_id)
            user_id = query.from_user.id
            task = self.tasks[user_id][task_id]
            
            if action == 'pause':
                task.status = 'paused'
                status_msg = "⏸️ Task paused"
            elif action == 'resume':
                task.status = 'running'
                status_msg = "▶️ Task resumed"
            else:  # stop
                task.status = 'stopped'
                status_msg = "⏹️ Task stopped"
            
            await query.message.edit_text(
                f"{status_msg}\n\n"
                f"Token: {task.token_address}\n"
                f"Bot: {task.username}",
                reply_markup=self.get_task_management_keyboard(task_id)
            )
            return CHOOSING

        elif query.data == 'stop_bots':
            user_id = query.from_user.id
            if user_id in self.active_bots:
                for bot in self.active_bots[user_id]:
                    bot['active'] = False
                del self.active_bots[user_id]
                await query.message.edit_text(
                    "✅ All bots have been stopped!\n\n"
                    "What would you like to do next?",
                    reply_markup=self.get_main_keyboard()
                )
            else:
                await query.message.edit_text(
                    "❌ No active bots found!\n\n"
                    "What would you like to do?",
                    reply_markup=self.get_main_keyboard()
                )
            return CHOOSING

    async def handle_channel_input(self, update, context):
        """Handle channel input"""
        channel_info = update.message.text.split('\n', 1)
        channel_name = channel_info[0].strip()
        description = channel_info[1].strip() if len(channel_info) > 1 else "Crypto trading signals and alpha"
        
        await update.message.reply_text(
            "🎯 Starting Telegram promotion campaign...\n"
            "Will promote across 20 different tokens.\n"
            "This may take some time."
        )
        
        # Start promotion in separate thread
        thread = threading.Thread(
            target=self.run_telegram_promo,
            args=(update.message, channel_name, description)
        )
        thread.start()
        
        return CHOOSING

    def run_telegram_promo(self, message, channel_name: str, description: str):
        """Run Telegram promotion campaign"""
        try:
            # Create bot instance
            bot = PumpBot(self.captcha_api_key, use_proxy=False)
            
            # Setup bot
            wallet = bot.create_wallet()
            bot.register()
            bot.login()
            
            # Generate cool username
            username = self.comment_generator.generate_username()
            bio = f""
            bot.setup_profile(username, bio)
                        
            # Start promotion
            tokens_commented = bot.promote_telegram_across_tokens(
                channel_name=channel_name,
                description=description,
                comment_generator=self.comment_generator,
                max_tokens=20
            )
            
            # Send summary to channel
            summary = (
                f"✅ Promotion Campaign Complete\n"
                f"• Comments Posted: {tokens_commented}\n"
                f"• Target Tokens: 20\n"
                f"• Success Rate: {(tokens_commented/20)*100:.1f}%"
            )
            self.send_telegram_message(channel_name, summary)
            
            asyncio.run(message.reply_text(
                f"✅ Telegram promotion completed!\n"
                f"Successfully promoted in {tokens_commented} tokens.\n"
                f"Summary sent to {channel_name}\n\n"
                "What would you like to do next?",
                reply_markup=self.get_main_keyboard()
            ))
            
        except Exception as e:
            asyncio.run(message.reply_text(
                f"❌ Error during promotion: {e}\n\n"
                "What would you like to do next?",
                reply_markup=self.get_main_keyboard()
            ))

    def generate_bullish_comments(self, token, quantity):
        """Generate bullish comments using AI"""
        return self.comment_generator.generate_multiple_comments(token, quantity)

    def run_multiple_bots(self, message, token_address: str, quantity: int):
        comments_made = 0
        chunk_message = ""
        user_id = message.chat.id
        
        try:
            # Get pool stats before starting
            stats = self.user_pool.get_pool_stats()
            print(f"Pool stats before run: {stats}")
            
            for i in range(quantity):
                # Force new user creation periodically to maintain pool
                force_new = (i % 5 == 0 and stats['total_users'] < self.user_pool.min_pool_size)
                
                # Try to get user from pool
                user = self.user_pool.get_available_user(force_new=force_new)
                
                try:
                    bot = PumpBot(self.captcha_api_key, use_proxy=False)
                    
                    if user:
                        # Reuse existing user
                        bot.wallet = user['wallet']
                        success = bot.login()
                        if not success:
                            raise Exception("Login failed")
                        username = user['username']
                    else:
                        # Create new user
                        wallet = bot.create_wallet()
                        bot.register()
                        bot.login()
                        username = self.comment_generator.generate_username()
                        bio = f"Crypto Degen | Following {token_address}"
                        bot.setup_profile(username, bio)
                        self.user_pool.add_user(wallet, username, bio)
                        user = self.user_pool.get_available_user()

                    # Post comment
                    comment = bot.monitor_token_with_ai_comments(
                        token_address, 
                        self.comment_generator,
                        single_comment=True
                    )
                    
                    if comment:
                        comments_made += 1
                        chunk_message += f"{comments_made}. [{username}] {comment} ✅\n"
                        self.user_pool.release_user(user, success=True)
                        
                        if comments_made % 10 == 0:
                            asyncio.run(message.reply_text(
                                f"💬 Comments posted ({comments_made}/{quantity}):\n\n{chunk_message}"
                            ))
                            chunk_message = ""
                    
                    time.sleep(random.uniform(5, 15))
                    
                except Exception as e:
                    print(f"Error with bot {i+1}: {e}")
                    if user:
                        self.user_pool.release_user(user, success=False, error=e)
                    continue

            # Show final pool stats
            stats = self.user_pool.get_pool_stats()
            print(f"Pool stats after run: {stats}")
            
            asyncio.run(message.reply_text(
                f"✅ Finished posting comments!\n"
                f"Successfully posted: {comments_made}/{quantity}\n\n"
                "What would you like to do next?",
                reply_markup=self.get_main_keyboard()
            ))
            
        except Exception as e:
            asyncio.run(message.reply_text(
                f"❌ Error posting comments: {e}\n\n"
                "What would you like to do next?",
                reply_markup=self.get_main_keyboard()
            ))

    async def handle_comment_token_input(self, update, context):
        """Handle token input for comment generation"""
        self.current_token = update.message.text
        await update.message.reply_text(
            f"🎯 How many comments would you like to generate for {self.current_token}?",
            reply_markup=self.get_comment_quantity_keyboard()
        )
        return CHOOSING

    async def run(self):
        """Start the Telegram bot"""
        app = Application.builder().token(self.telegram_token).build()

        # Create conversation handler with simplified states
        conv_handler = ConversationHandler(
            entry_points=[CommandHandler('start', self.start)],
            states={
                CHOOSING: [
                    CallbackQueryHandler(self.button_handler)
                ],
                INPUT_CHANNEL: [
                    MessageHandler(filters.TEXT & ~filters.COMMAND, self.handle_channel_input),
                    CallbackQueryHandler(self.button_handler)
                ],
                GENERATE_COMMENTS: [
                    MessageHandler(filters.TEXT & ~filters.COMMAND, self.handle_comment_token_input),
                    CallbackQueryHandler(self.button_handler)
                ]
            },
            fallbacks=[CommandHandler('start', self.start)]
        )

        app.add_handler(conv_handler)

        # Start the bot
        print("Starting bot...")
        await app.initialize()
        await app.start()

        try:
            await app.updater.start_polling(allowed_updates=Update.ALL_TYPES)
            print("Bot is running! Press Ctrl+C to stop.")
            
            # Keep the bot running
            stop_event = asyncio.Event()
            while not stop_event.is_set():
                try:
                    await asyncio.sleep(1)
                except asyncio.CancelledError:
                    stop_event.set()
                    break
                
        except Exception as e:
            print(f"Error while running bot: {e}")
        finally:
            print("Stopping bot...")
            await app.stop()

    def send_telegram_message(self, channel_name: str, message: str) -> bool:
        """Send message to a Telegram channel"""
        try:
            # Make sure channel name starts with @
            if not channel_name.startswith('@'):
                channel_name = '@' + channel_name
            
            # Send message using Telegram Bot API
            url = f"https://api.telegram.org/bot{self.telegram_token}/sendMessage"
            data = {
                'chat_id': channel_name,
                'text': message,
                'parse_mode': 'HTML'  # Allow basic formatting
            }
            
            response = requests.post(url, json=data)
            result = response.json()
            
            if result.get('ok'):
                print(f"✅ Message sent to {channel_name}")
                return True
            else:
                print(f"❌ Failed to send message: {result.get('description')}")
                return False
            
        except Exception as e:
            print(f"Error sending Telegram message: {e}")
            return False

if __name__ == '__main__':
    # Create and run bot
    bot_manager = TelegramManager()
    try:
        asyncio.run(bot_manager.run())
    except KeyboardInterrupt:
        print("\nBot stopped by user")
    except Exception as e:
        print(f"\nError running bot: {e}")