main.py

"""
CleanAcc - Enhanced Telegram Account Cleaner Bot
Repository: https://github.com/MeherMankar/CleanAcc
Developer: @MeherMankar (https://github.com/MeherMankar)
"""

import asyncio
import logging
import os
import sys
from datetime import datetime
from aiogram import Bot, Dispatcher
from aiogram.enums import ParseMode
import redis.asyncio as redis
from contextlib import asynccontextmanager

import config
from database import init_db, db_pool
from handlers import start, auth, cabinet, cleanup, admin
from services.proxy_manager import ProxyManager

# Global variables for service access
global_proxy_manager = None

def setup_logging():
    """Setup structured logging with file and console output"""
    # Create logs directory if it doesn't exist
    os.makedirs('logs', exist_ok=True)
    
    # Configure logging format
    log_format = '%(asctime)s - %(name)s - %(levelname)s - %(funcName)s:%(lineno)d - %(message)s'
    
    # Setup file handler with rotation
    file_handler = logging.FileHandler(
        f'logs/{config.LOG_FILE}',
        encoding='utf-8'
    )
    file_handler.setLevel(getattr(logging, config.LOG_LEVEL))
    file_handler.setFormatter(logging.Formatter(log_format))
    
    # Setup console handler
    console_handler = logging.StreamHandler(sys.stdout)
    console_handler.setLevel(logging.INFO)
    console_handler.setFormatter(logging.Formatter(
        '%(asctime)s - %(levelname)s - %(message)s'
    ))
    
    # Configure root logger
    logging.basicConfig(
        level=getattr(logging, config.LOG_LEVEL),
        handlers=[file_handler, console_handler],
        format=log_format
    )
    
    # Reduce noise from external libraries
    logging.getLogger('aiohttp').setLevel(logging.WARNING)
    logging.getLogger('aiogram').setLevel(logging.INFO)
    logging.getLogger('telethon').setLevel(logging.WARNING)
    
    logger = logging.getLogger(__name__)
    logger.info(f"Logging initialized - Level: {config.LOG_LEVEL}")
    return logger

@asynccontextmanager
async def lifespan_manager():
    """Manage application lifespan with proper cleanup"""
    logger = logging.getLogger(__name__)
    
    # Startup
    logger.info("Starting CleanAcc Bot...")
    
    # Initialize Redis with password if provided
    redis_client = redis.Redis(
        host=config.REDIS_HOST,
        port=config.REDIS_PORT,
        password=config.REDIS_PASSWORD if config.REDIS_PASSWORD else None,
        decode_responses=True,
        retry_on_timeout=True,
        health_check_interval=30
    )
    
    # Test Redis connection
    try:
        await redis_client.ping()
        logger.info("Redis connection established")
    except Exception as e:
        logger.warning(f"Redis connection failed: {e}")
        logger.warning("Continuing without Redis - some features may be limited")
        redis_client = None
    
    # Initialize database pool
    await db_pool.initialize()
    await init_db()
    
    # Initialize proxy manager
    proxy_manager = None
    if redis_client:
        proxy_manager = ProxyManager(redis_client)
        await proxy_manager.start_background_update()
    else:
        logger.warning("Proxy manager disabled - Redis not available")
    
    global global_proxy_manager
    global_proxy_manager = proxy_manager
    
    try:
        yield {
            'redis_client': redis_client,
            'proxy_manager': proxy_manager
        }
    finally:
        # Cleanup
        logger.info("Shutting down CleanAcc Bot...")
        
        try:
            if proxy_manager:
                await proxy_manager.stop_background_update()
            if redis_client:
                await redis_client.close()
            await db_pool.close()
            logger.info("Cleanup completed successfully")
        except Exception as e:
            logger.error(f"Error during cleanup: {e}")


async def health_check():
    """Perform health checks on all services"""
    logger = logging.getLogger(__name__)
    health_status = {
        'timestamp': datetime.now().isoformat(),
        'services': {}
    }
    
    # Check database
    try:
        async with db_pool.get_connection() as db:
            await db.execute("SELECT 1")
        health_status['services']['database'] = 'healthy'
    except Exception as e:
        health_status['services']['database'] = f'unhealthy: {e}'
        logger.error(f"Database health check failed: {e}")
    
    # Check proxy manager
    try:
        if global_proxy_manager:
            stats = await global_proxy_manager.get_proxy_stats()
            health_status['services']['proxy_manager'] = {
                'status': 'healthy',
                'stats': stats
            }
        else:
            health_status['services']['proxy_manager'] = 'not_initialized'
    except Exception as e:
        health_status['services']['proxy_manager'] = f'unhealthy: {e}'
        logger.error(f"Proxy manager health check failed: {e}")
    
    return health_status

async def main():
    """Enhanced main function with proper error handling and monitoring"""
    logger = setup_logging()
    
    try:
        # Validate configuration
        logger.info("Validating configuration...")
        config.validate_config()
        
        async with lifespan_manager() as services:
            redis_client = services['redis_client']
            proxy_manager = services['proxy_manager']
            
            # Create sessions folder if it doesn't exist
            os.makedirs(config.SESSIONS_DIR, exist_ok=True)
            logger.info(f"Sessions directory: {config.SESSIONS_DIR}")
            
            # Bot initialization with error handling
            try:
                bot = Bot(
                    token=config.BOT_TOKEN,
                    parse_mode=ParseMode.HTML
                )
                
                # Test bot token
                bot_info = await bot.get_me()
                logger.info(f"Bot initialized: @{bot_info.username} ({bot_info.first_name})")
                
            except Exception as e:
                logger.error(f"Bot initialization failed: {e}")
                raise
            
            # Dispatcher initialization
            dp = Dispatcher()
            
            # Add middleware for rate limiting and logging
            from middlewares.rate_limit import RateLimitMiddleware
            from middlewares.logging import LoggingMiddleware
            
            dp.message.middleware(RateLimitMiddleware())
            dp.callback_query.middleware(RateLimitMiddleware())
            dp.message.middleware(LoggingMiddleware())
            dp.callback_query.middleware(LoggingMiddleware())
            
            # Handler registration - use enhanced handlers
            dp.include_router(start.router)
            
            dp.include_router(auth.router)
            dp.include_router(cabinet.router)
            dp.include_router(cleanup.router)
            
            dp.include_router(admin.router)
            
            # Start health check task
            health_task = asyncio.create_task(periodic_health_check())
            
            try:
                logger.info("Starting bot polling...")
                await dp.start_polling(
                    bot,
                    allowed_updates=['message', 'callback_query'],
                    drop_pending_updates=True
                )
            finally:
                health_task.cancel()
                try:
                    await health_task
                except asyncio.CancelledError:
                    pass
                
                await bot.session.close()
                logger.info("Bot session closed")
    
    except KeyboardInterrupt:
        logger.info("Received shutdown signal")
    except Exception as e:
        logger.error(f"Fatal error: {e}", exc_info=True)
        raise
    finally:
        logger.info("Application shutdown complete")

async def periodic_health_check():
    """Periodic health check task"""
    logger = logging.getLogger(__name__)
    
    while True:
        try:
            await asyncio.sleep(300)  # Check every 5 minutes
            health_status = await health_check()
            
            # Log health status
            unhealthy_services = [
                service for service, status in health_status['services'].items()
                if isinstance(status, str) and 'unhealthy' in status
            ]
            
            if unhealthy_services:
                logger.warning(f"Unhealthy services detected: {unhealthy_services}")
            else:
                logger.debug("All services healthy")
                
        except asyncio.CancelledError:
            break
        except Exception as e:
            logger.error(f"Health check error: {e}")


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\nShutdown requested by user")
    except Exception as e:
        print(f"Fatal error: {e}")
        sys.exit(1)