Update core/bot.py

core/bot.py

@@ -1,4 +1,4 @@
-# filename: core/bot.py
+# filename: core/bot.py (Corrected Version)
 
 import asyncio
 import logging
@@ -14,8 +14,10 @@ from utils import terabox, ffmpeg, helpers
 
 logger = logging.getLogger(__name__)
 
-# --- Core Bot Client ---
-bot = TelegramClient('terabox_bot_session', config.API_ID, config.API_HASH)
+# --- Core Bot Client (THIS LINE HAS BEEN CHANGED) ---
+# We now specify the 'data/' directory for the session file to ensure it's writable.
+bot = TelegramClient('data/terabox_bot_session', config.API_ID, config.API_HASH)
+
 
 # --- In-Memory State Management ---
 # Queues for the workers
@@ -58,10 +60,11 @@ async def scheduler_loop():
             
             # If the user has no more tasks, remove them from the turn order.
             # Otherwise, move them to the back of the line.
-            USER_TURN_ORDER.rotate(-1) # Move current user to the end
             if not ACTIVE_USER_TASKS[user_id]:
-                USER_TURN_ORDER.pop() # Remove them if their list is now empty
+                USER_TURN_ORDER.popleft() # Remove them if their list is now empty
                 del ACTIVE_USER_TASKS[user_id]
+            else:
+                USER_TURN_ORDER.rotate(-1) # Move current user to the end
         else:
             # Cleanup: If user is in the turn order but has no tasks, remove them.
             USER_TURN_ORDER.popleft()
@@ -73,119 +76,146 @@ async def _process_task(task: dict, worker_name: str):
     batch_id = task['batch_id']
     original_link = task['link']
     user_id = task['user_id']
-    metadata = task['metadata'] # Metadata is pre-fetched by the handler
+    metadata = task['metadata']
+    is_cached_task = task.get('cached', False)
     
     batch_info = BATCH_JOBS.get(batch_id)
     if not batch_info:
         logger.warning(f"[{worker_name}] Batch {batch_id} not found. Task for {original_link} skipped.")
         return
 
-    logger.info(f"[{worker_name}] Starting processing for link: {original_link}")
+    logger.info(f"[{worker_name}] Starting processing for link: {original_link} (Cached: {is_cached_task})")
     
     download_path = None
     final_file_path = None
+    thumbnail_path = None
     error_reason = None
-
+    
     try:
-        # Step 1: Download the file from the direct link
-        download_path = await terabox.download_file_from_url(
-            url=metadata['url'],
-            dir_path="downloads",
-            filename=metadata['file_name']
-        )
-        if not download_path:
-            raise ValueError("File download failed.")
-        
-        final_file_path = download_path
-        
-        # Step 2: FFMPEG processing (if enabled and it's a video)
-        thumbnail_path = None
-        if config.ENABLE_FFMPEG and final_file_path.endswith(('.mp4', '.mkv', '.webm')):
-            # Remux to MP4 if needed (fast operation)
-            if not final_file_path.endswith('.mp4'):
-                remuxed_path = f"{os.path.splitext(final_file_path)[0]}.mp4"
-                remuxed_path = await ffmpeg.remux_to_mp4(final_file_path, remuxed_path)
-                if remuxed_path:
-                    os.remove(final_file_path) # remove original file
-                    final_file_path = remuxed_path
+        if is_cached_task:
+            # For cached tasks, we just need to forward the file
+            cached_file_id = metadata['file_id']
+            dm_message = await bot.send_file(
+                user_id,
+                file=int(cached_file_id), # Use file_id from cache
+                caption=templates.BotResponses.FILE_CAPTION.format(
+                    file_name=metadata['file_name'],
+                    file_size=helpers.format_bytes(metadata['file_size']),
+                    source_url=original_link
+                )
+            )
+        else:
+            # For new tasks, follow the full download -> process -> upload workflow
+            download_path = await terabox.download_file_from_url(
+                url=metadata['url'],
+                dir_path="downloads",
+                filename=metadata['file_name']
+            )
+            if not download_path:
+                raise ValueError("File download failed.")
             
-            # Generate thumbnail (fast operation)
-            thumb_path = f"{os.path.splitext(final_file_path)[0]}.jpg"
-            thumbnail_path = await ffmpeg.generate_thumbnail(final_file_path, thumb_path)
-
-        # Step 3: Upload to backup channel & cache
-        caption = templates.BotResponses.FILE_CAPTION.format(
-            file_name=metadata['file_name'],
-            file_size=helpers.format_bytes(metadata['file_size']),
-            source_url=original_link
-        )
-        
-        backup_message = await bot.send_file(
-            config.BACKUP_CHANNEL_ID,
-            file=final_file_path,
-            thumb=thumbnail_path,
-            caption=caption
-        )
-        await db_manager.add_to_cache(
-            short_id=task['short_id'],
-            file_id=backup_message.id,
-            file_name=metadata['file_name'],
-            file_size=metadata['file_size']
-        )
-
-        # Step 4: Deliver to user and schedule deletion
-        dm_message = await bot.send_file(
-            user_id,
-            file=backup_message, # Send using file_id from backup channel for speed
-            caption=caption
-        )
-        # Here you would add logic to call the APScheduler to delete dm_message in 30 mins
+            final_file_path = download_path
+            
+            if config.ENABLE_FFMPEG and final_file_path.lower().endswith(('.mp4', '.mkv', '.webm')):
+                if not final_file_path.lower().endswith('.mp4'):
+                    remuxed_path = f"{os.path.splitext(final_file_path)[0]}.mp4"
+                    remuxed_path = await ffmpeg.remux_to_mp4(final_file_path, remuxed_path)
+                    if remuxed_path:
+                        if os.path.exists(final_file_path): os.remove(final_file_path)
+                        final_file_path = remuxed_path
+                
+                thumb_path_temp = f"{os.path.splitext(final_file_path)[0]}.jpg"
+                thumbnail_path = await ffmpeg.generate_thumbnail(final_file_path, thumb_path_temp)
+
+            caption = templates.BotResponses.FILE_CAPTION.format(
+                file_name=metadata['file_name'],
+                file_size=helpers.format_bytes(metadata['file_size']),
+                source_url=original_link
+            )
+            
+            backup_message = await bot.send_file(
+                config.BACKUP_CHANNEL_ID,
+                file=final_file_path,
+                thumb=thumbnail_path,
+                caption=caption
+            )
+            await db_manager.add_to_cache(
+                short_id=task['short_id'],
+                file_id=str(backup_message.id),
+                file_name=metadata['file_name'],
+                file_size=metadata['file_size']
+            )
 
+            dm_message = await bot.send_file(
+                user_id,
+                file=backup_message,
+                caption=caption
+            )
     except Exception as e:
         logger.error(f"[{worker_name}] Error processing {original_link}: {e}", exc_info=True)
         error_reason = str(e)
 
     finally:
-        # Clean up local files
-        if download_path and os.path.exists(download_path):
-            os.remove(download_path)
-        if final_file_path and final_file_path != download_path and os.path.exists(final_file_path):
-            os.remove(final_file_path)
-        if 'thumbnail_path' in locals() and thumbnail_path and os.path.exists(thumbnail_path):
-            os.remove(thumbnail_path)
-    
+        # Clean up local temporary files
+        for path in [download_path, final_file_path, thumbnail_path]:
+            if path and os.path.exists(path):
+                try:
+                    os.remove(path)
+                except OSError as e:
+                    logger.error(f"Error removing file {path}: {e}")
+
     # --- Final Step: Update Batch Status Safely ---
     async with batch_info['lock']:
         batch_info['processed_links'] += 1
         if error_reason:
             batch_info['failed_links'].append({"link": original_link, "error": error_reason})
         
-        # Update progress message
         processed = batch_info['processed_links']
         total = batch_info['total_links']
         progress = processed / total
         
         try:
+            # Update progress message
             await bot.edit_message(
                 batch_info['chat_id'],
                 batch_info['status_message_id'],
                 text=templates.BotResponses.BATCH_UPDATE_VALIDATED.format(
-                    batch_id=batch_id[:6],
+                    batch_id=batch_id,
                     valid_count=total,
                     total_count=batch_info['original_total'],
-                    skipped_count=batch_info['original_total'] - total,
+                    skipped_count=len(batch_info['skipped_links']),
                     progress_bar=helpers.create_progress_bar(progress),
                     processed_count=processed
                 )
             )
         except Exception:
-            pass # Ignore if message can't be edited
+            pass 
 
         # If batch is fully processed, send final summary
         if processed == total:
-            # Here you would call a final summary function
-            logger.info(f"[{worker_name}] Batch {batch_id[:6]} complete.")
-            # ... logic to send final summary and delete job from BATCH_JOBS ...
+            logger.info(f"[{worker_name}] Batch {batch_id} complete.")
+            failed_items = batch_info['skipped_links'] + batch_info['failed_links']
+            
+            if failed_items:
+                failed_details = "\n".join([f"- `{item['link']}` ({item['error']})" for item in failed_items])
+                final_text = templates.BotResponses.BATCH_COMPLETE_SUMMARY.format(
+                    batch_id=batch_id,
+                    success_count=total - len(failed_items),
+                    skipped_count=len(failed_items),
+                    failed_details=failed_details
+                )
+            else:
+                final_text = templates.BotResponses.BATCH_COMPLETE_NO_SKIPS.format(
+                    batch_id=batch_id,
+                    success_count=total
+                )
+            
+            try:
+                await bot.edit_message(batch_info['chat_id'], batch_info['status_message_id'], final_text)
+            except Exception:
+                pass
+            
+            del BATCH_JOBS[batch_id]
 
 
 async def worker(name: str, queue: asyncio.Queue):
@@ -197,19 +227,15 @@ async def worker(name: str, queue: asyncio.Queue):
         finally:
             queue.task_done()
 
-# --- Startup Function ---
+
 def start_bot_runtime():
     """Creates all the background tasks for the bot's engine."""
-    # Create Premium Workers
     for i in range(config.PREMIUM_WORKERS):
         asyncio.create_task(worker(f"PremiumWorker-{i+1}", PREMIUM_QUEUE))
     
-    # Create Free Workers
     for i in range(config.FREE_WORKERS):
         asyncio.create_task(worker(f"FreeWorker-{i+1}", FREE_QUEUE))
     
-    # Start the Fair-Share Scheduler
     asyncio.create_task(scheduler_loop())
     
     logger.info(f"Bot runtime started with {config.PREMIUM_WORKERS} premium and {config.FREE_WORKERS} free workers.")
-