Create links.py

handlers/links.py

@@ -0,0 +1,145 @@
+# filename: handlers/links.py
+
+import logging
+import re
+import uuid
+import asyncio
+from telethon import events, Button
+
+from core.bot import bot, BATCH_JOBS, PREMIUM_QUEUE, FREE_QUEUE, ACTIVE_USER_TASKS, USER_TURN_ORDER
+import config
+import templates
+from utils import terabox, helpers
+from database import manager as db_manager
+from datetime import datetime
+
+logger = logging.getLogger(__name__)
+
+# This regular expression finds all URLs in a message
+URL_REGEX = r'https?://[^\s<>"\']+'
+
+@bot.on(events.NewMessage(func=lambda e: e.text and not e.text.startswith('/')))
+async def main_link_handler(event):
+    sender = await event.get_sender()
+    user_id = sender.id
+    
+    # --- Step 1: Pre-flight Checks ---
+    user = await db_manager.get_user(user_id)
+    if not user:
+        user = await db_manager.add_or_update_user(user_id, sender.username, sender.first_name)
+    
+    if user.get('is_banned'):
+        await event.reply(templates.BotResponses.USER_BANNED_MESSAGE)
+        return
+
+    if event.is_group and not await db_manager.is_group_authorized(event.chat_id):
+        # Only reply if the group is not authorized, to avoid spam
+        # You could add a check to only reply once per hour per group
+        # await event.reply(templates.BotResponses.GROUP_NOT_AUTHORIZED.format(owner_id=config.OWNER_ID))
+        return
+
+    # Force Subscribe Check
+    if config.FORCE_SUB_CHANNEL_USERNAME:
+        try:
+            await bot.get_permissions(config.FORCE_SUB_CHANNEL_USERNAME, user_id)
+        except Exception:
+            await event.reply(templates.BotResponses.FORCE_SUBSCRIBE_MESSAGE.format(channel_username=config.FORCE_SUB_CHANNEL_USERNAME))
+            return
+
+    # --- Step 2: Parse Links and Create Batch ---
+    links = list(set(re.findall(URL_REGEX, event.text)))
+    terabox_links = [link for link in links if "terabox" in link or "terashare" in link]
+
+    if not terabox_links:
+        return
+
+    batch_id = str(uuid.uuid4())[:6]
+    status_msg = await event.reply(templates.BotResponses.BATCH_ACKNOWLEDGEMENT.format(
+        link_count=len(terabox_links),
+        batch_id=batch_id
+    ))
+
+    # --- Step 3: Evaluate Each Link ---
+    valid_tasks = []
+    skipped_links = []
+
+    for link in terabox_links:
+        short_id = await terabox.extract_terabox_short_id(link)
+        if not short_id:
+            skipped_links.append({"link": link, "error": "Invalid Link Format"})
+            continue
+
+        cached = await db_manager.get_cached_file(short_id)
+        task_data = {"batch_id": batch_id, "link": link, "user_id": user_id, "short_id": short_id}
+        
+        if cached:
+            task_data['metadata'] = {"file_name": cached['file_name'], "file_size": cached['file_size']}
+            task_data['cached'] = True
+            valid_tasks.append(task_data)
+            continue
+            
+        # If not cached, get metadata for size check
+        metadata = await terabox.get_final_url_and_metadata(link)
+        if not metadata['success']:
+            skipped_links.append({"link": link, "error": metadata['error']})
+            continue
+
+        task_data['metadata'] = metadata
+        task_data['cached'] = False
+        
+        # Apply free user limit
+        is_premium = user.get('is_premium') and user.get('premium_expiry_date', datetime.min) > datetime.utcnow()
+        if not is_premium and metadata['file_size'] > config.FREE_USER_FILE_SIZE_LIMIT_BYTES:
+            skipped_links.append({
+                "link": link,
+                "error": templates.BotResponses.PREMIUM_REQUIRED_ERROR.format(
+                    file_name=metadata['file_name'],
+                    file_size=helpers.format_bytes(metadata['file_size']),
+                    free_limit=helpers.format_bytes(config.FREE_USER_FILE_SIZE_LIMIT_BYTES)
+                )
+            })
+            continue
+        
+        valid_tasks.append(task_data)
+
+    # --- Step 4: Queue the Valid Tasks ---
+    if not valid_tasks:
+        # Handle case where no links were valid
+        await status_msg.edit("❌ All links provided were invalid or failed the initial check.")
+        return
+
+    # Update the batch job tracker
+    BATCH_JOBS[batch_id] = {
+        "user_id": user_id,
+        "chat_id": event.chat_id,
+        "status_message_id": status_msg.id,
+        "total_links": len(valid_tasks),
+        "processed_links": 0,
+        "original_total": len(terabox_links),
+        "failed_links": skipped_links, # pre-flight failures
+        "lock": asyncio.Lock()
+    }
+    
+    is_premium = user.get('is_premium') and user.get('premium_expiry_date', datetime.min) > datetime.utcnow()
+    
+    if is_premium:
+        for task in valid_tasks:
+            await PREMIUM_QUEUE.put(task)
+    else:
+        # For free users, add tasks to their personal list and add them to the round-robin deque
+        if user_id not in ACTIVE_USER_TASKS:
+            ACTIVE_USER_TASKS[user_id] = []
+        ACTIVE_USER_TASKS[user_id].extend(valid_tasks)
+        
+        if user_id not in USER_TURN_ORDER:
+            USER_TURN_ORDER.append(user_id)
+
+    # Update the status message to show how many links are being processed
+    await status_msg.edit(templates.BotResponses.BATCH_UPDATE_VALIDATED.format(
+        batch_id=batch_id,
+        valid_count=len(valid_tasks),
+        total_count=len(terabox_links),
+        skipped_count=len(skipped_links),
+        progress_bar=helpers.create_progress_bar(0),
+        processed_count=0
+    ))