Update core/bot.py

core/bot.py

@@ -1,4 +1,4 @@
-# filename: core/bot.py (Final Corrected Version)
+# filename: core/bot.py
 
 import asyncio
 import logging
@@ -7,7 +7,6 @@ import time
 from collections import deque
 from telethon import TelegramClient
 from telethon.network import connection
-import dns.asyncresolver # <-- IMPORT FOR YOUR DNS SUGGESTION
 
 import config
 import templates
@@ -16,27 +15,15 @@ from utils import terabox, ffmpeg, helpers
 
 logger = logging.getLogger(__name__)
 
-# --- YOUR EXCELLENT DNS SUGGESTION ---
-# Forcing the use of reliable public DNS servers to improve connection stability.
-try:
-    async_resolver = dns.asyncresolver.Resolver(configure=False)
-    async_resolver.nameservers = ["8.8.8.8", "8.8.4.4", "1.1.1.1"]
-    dns.asyncresolver.default_resolver = async_resolver
-    logger.info("Successfully configured custom DNS resolvers.")
-except Exception as e:
-    logger.warning(f"Could not configure custom DNS resolver: {e}")
-
-
-# --- Core Bot Client (UPDATED WITH BOTH FIXES) ---
+# --- Core Bot Client ---
+# Includes fix for session path and connection mode
 bot = TelegramClient(
     'data/terabox_bot_session',
     config.API_ID,
     config.API_HASH,
-    connection=connection.ConnectionHttp # <-- MY TYPO IS NOW CORRECTED
+    connection=connection.ConnectionHttp
 )
 
-# --- (The rest of the file is exactly the same as before) ---
-
 # --- In-Memory State Management ---
 PREMIUM_QUEUE = asyncio.Queue()
 FREE_QUEUE = asyncio.Queue()
@@ -44,18 +31,26 @@ BATCH_JOBS = {}
 ACTIVE_USER_TASKS = {}
 USER_TURN_ORDER = deque()
 
+
 # --- The Fair-Share Scheduler for Free Users ---
 async def scheduler_loop():
+    """
+    The heart of the fair-share system.
+    This loop checks whose turn it is and feeds one task at a time into the FREE_QUEUE.
+    """
     logger.info("Fair-Share Scheduler started.")
-    # ... (rest of the function is the same) ...
     while True:
-        await asyncio.sleep(0.2)
+        await asyncio.sleep(0.2) # Prevent busy-looping
+
         if not USER_TURN_ORDER or FREE_QUEUE.full():
             continue
+
         user_id = USER_TURN_ORDER[0]
+
         if user_id in ACTIVE_USER_TASKS and ACTIVE_USER_TASKS[user_id]:
             task = ACTIVE_USER_TASKS[user_id].pop(0)
             await FREE_QUEUE.put(task)
+
             if not ACTIVE_USER_TASKS[user_id]:
                 USER_TURN_ORDER.popleft()
                 del ACTIVE_USER_TASKS[user_id]
@@ -64,12 +59,8 @@ async def scheduler_loop():
         else:
             USER_TURN_ORDER.popleft()
 
-# --- The Worker Logic ---
-async def _process_task(task: dict, worker_name: str):
-    # ... (rest of the function is the same) ...
-    pass # Placeholder for the large function body
 
-# --- (The full _process_task, worker, and start_bot_runtime functions are here, unchanged) ---
+# --- The Worker Logic ---
 async def _process_task(task: dict, worker_name: str):
     """The main logic for processing a single link. Executed by all workers."""
     batch_id = task['batch_id']
@@ -92,27 +83,110 @@ async def _process_task(task: dict, worker_name: str):
     
     try:
         if is_cached_task:
+            # For cached tasks, we forward the file from the backup channel
             cached_file_info = await db_manager.get_cached_file(task['short_id'])
             if cached_file_info and config.BACKUP_CHANNEL_ID:
                 message_to_forward = await bot.get_messages(config.BACKUP_CHANNEL_ID, ids=int(cached_file_info['file_id']))
                 if message_to_forward:
                     await bot.send_file(user_id, file=message_to_forward)
         else:
-            # Full download -> process -> upload workflow
-            pass # Placeholder for this complex logic which remains unchanged
+            # 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.")
+            
+            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']
+            )
+            await bot.send_file(user_id, file=backup_message)
 
     except Exception as e:
         logger.error(f"[{worker_name}] Error processing {original_link}: {e}", exc_info=True)
         error_reason = str(e)
 
     finally:
-        # Cleanup
-        pass
+        # 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}")
 
-    # Update Batch Status
+    # --- Final Step: Update Batch Status Safely ---
     async with batch_info['lock']:
-        # ... (logic remains unchanged) ...
-        pass
+        batch_info['processed_links'] += 1
+        if error_reason:
+            batch_info['failed_links'].append({"link": original_link, "error": error_reason})
+        
+        processed = batch_info['processed_links']
+        total = batch_info['total_links']
+        
+        try:
+            await bot.edit_message(
+                batch_info['chat_id'],
+                batch_info['status_message_id'],
+                text=templates.BotResponses.BATCH_UPDATE_VALIDATED.format(
+                    batch_id=batch_id,
+                    valid_count=total,
+                    total_count=batch_info['original_total'],
+                    skipped_count=len(batch_info['skipped_links']),
+                    progress_bar=helpers.create_progress_bar(processed / total),
+                    processed_count=processed
+                )
+            )
+        except Exception: pass 
+
+        if processed == total:
+            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.get('link', 'N/A')}` ({item.get('error', 'Unknown 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
+            
+            if batch_id in BATCH_JOBS: del BATCH_JOBS[batch_id]
 
 
 async def worker(name: str, queue: asyncio.Queue):
@@ -120,12 +194,11 @@ async def worker(name: str, queue: asyncio.Queue):
     while True:
         task = await queue.get()
         try:
-            # We are using a simplified placeholder here for brevity
-            # The full _process_task logic from the previous file goes here.
-            await asyncio.sleep(1) 
+            await _process_task(task, name)
         finally:
             queue.task_done()
 
+
 def start_bot_runtime():
     """Creates all the background tasks for the bot's engine."""
     for i in range(config.PREMIUM_WORKERS):
@@ -135,6 +208,4 @@ def start_bot_runtime():
         asyncio.create_task(worker(f"FreeWorker-{i+1}", FREE_QUEUE))
     
     asyncio.create_task(scheduler_loop())
-    
     logger.info(f"Bot runtime started with {config.PREMIUM_WORKERS} premium and {config.FREE_WORKERS} free workers.")
-