Update app.py

app.py

@@ -1,129 +1,205 @@
 import os
-import time
 import re
+import sys
+import time
 import base64
 import hashlib
 import shutil
-import threading
-import asyncio
+import logging
+import secrets
 import sqlite3
 import tempfile
-import secrets
-from typing import Optional
+import threading
+import asyncio
+from pathlib import Path
+from typing import Optional, List, Tuple, Union
 
 import requests
 from flask import Flask
 from pyrogram import Client, filters
-from pyrogram.types import InlineKeyboardMarkup, InlineKeyboardButton
-from huggingface_hub import HfApi, get_hf_file_metadata, hf_hub_download, hf_hub_url
+from pyrogram.types import (
+    InlineKeyboardMarkup,
+    InlineKeyboardButton,
+    Message,
+    CallbackQuery,
+)
+from huggingface_hub import (
+    HfApi,
+    get_hf_file_metadata,
+    hf_hub_download,
+    hf_hub_url,
+)
 from cryptography.hazmat.primitives.ciphers.aead import AESSIV
 
-# --------------- CONFIG ---------------
-API_ID = int(os.environ["API_ID"])
-API_HASH = os.environ["API_HASH"]
-BOT_TOKEN = os.environ["BOT_TOKEN"]
-HF_TOKEN = os.environ["HF_TOKEN"]
-OBFUSCATION_KEY = os.environ.get("OBFUSCATION_KEY", "change-me-to-a-secret-string")
+# =============================================================================
+# CONFIGURATION
+# =============================================================================
+
+API_ID: int = int(os.environ["API_ID"])
+API_HASH: str = os.environ["API_HASH"]
+BOT_TOKEN: str = os.environ["BOT_TOKEN"]
+HF_TOKEN: str = os.environ["HF_TOKEN"]
+OBFUSCATION_KEY: str = os.environ.get("OBFUSCATION_KEY", "change-me-to-a-secret-string")
+
+# Hard‑coded Space URL for keep‑alive pings
+SPACE_HOST: str = "https://cxops-push.hf.space"
+
+DATASET_REPO: str = "Cxops/Bzxv"
+DATASET_TYPE: str = "dataset"
+MAX_FILE_SIZE_TG: int = 1900 * 1024 * 1024      # 1900 MB per part
+MIN_FREE_DISK: int = 45 * 1024 * 1024 * 1024    # 45 GB
+DATASET_MAX_SIZE: int = int(8.7 * 1024**4)       # 8.7 TB
+KEEPALIVE_INTERVAL: int = 300                    # 5 minutes
+DB_PATH: str = "bot_files.db"
+
+# Logging
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+    handlers=[logging.StreamHandler(sys.stdout)],
+)
+logger = logging.getLogger("StorageBot")
+
+# =============================================================================
+# DETERMINISTIC ENCRYPTION (AES‑SIV)
+# =============================================================================
 
-DATASET_REPO = "Cxops/Bzxv"
-MAX_FILE_SIZE_TG = 1900 * 1024 * 1024      # 1900 MB per Telegram part
-MIN_FREE_DISK = 45 * 1024 * 1024 * 1024    # 45 GB
-DATASET_MAX_SIZE = 8.7 * 1024 * 1024 * 1024 * 1024  # 8.7 TB
-KEEPALIVE_INTERVAL = 300                   # 5 minutes
-DB_PATH = "bot_files.db"
+_siv_key = hashlib.sha512(OBFUSCATION_KEY.encode()).digest()
+_aes_siv = AESSIV(_siv_key)
 
-# --------------- DETERMINISTIC ENCRYPTION (AES‑SIV) ---------------
-siv_key = hashlib.sha512(OBFUSCATION_KEY.encode()).digest()
-aes_siv = AESSIV(siv_key)
 
 def encrypt_text(plain: str) -> str:
-    token = aes_siv.encrypt(plain.encode(), None)
+    """Return URL‑safe base64 (no padding) deterministic ciphertext."""
+    token = _aes_siv.encrypt(plain.encode(), None)
     return base64.urlsafe_b64encode(token).decode().rstrip("=")
 
+
 def decrypt_text(encoded: str) -> str:
+    """Recover original string from URL‑safe base64 (with or without padding)."""
     padding = 4 - len(encoded) % 4
     if padding != 4:
         encoded += "=" * padding
     token = base64.urlsafe_b64decode(encoded.encode())
-    return aes_siv.decrypt(token, None).decode()
+    return _aes_siv.decrypt(token, None).decode()
+
 
 def get_encrypted_folder(user_id: int) -> str:
     return encrypt_text(str(user_id))
 
+
 def get_encrypted_filename(original: str) -> str:
     return encrypt_text(original)
 
-# --------------- SQLITE DATABASE ---------------
+
+# =============================================================================
+# DATABASE
+# =============================================================================
+
 conn = sqlite3.connect(DB_PATH, check_same_thread=False)
+conn.row_factory = sqlite3.Row
+conn.execute("PRAGMA journal_mode=WAL")
 cursor = conn.cursor()
-cursor.execute("""
+cursor.execute(
+    """
     CREATE TABLE IF NOT EXISTS files (
-        user_id INTEGER,
-        enc_path TEXT UNIQUE,
+        user_id   INTEGER,
+        enc_path  TEXT UNIQUE,
         orig_name TEXT,
-        size INTEGER,
+        size      INTEGER,
         timestamp REAL
     )
-""")
+    """
+)
 conn.commit()
 
-def db_add_file(user_id, enc_path, orig_name, size):
+
+def db_add_file(user_id: int, enc_path: str, orig_name: str, size: int) -> None:
     cursor.execute(
         "INSERT OR REPLACE INTO files (user_id, enc_path, orig_name, size, timestamp) VALUES (?, ?, ?, ?, ?)",
-        (user_id, enc_path, orig_name, size, time.time())
+        (user_id, enc_path, orig_name, size, time.time()),
     )
     conn.commit()
 
-def db_get_user_files(user_id):
-    cursor.execute("SELECT enc_path, orig_name, size FROM files WHERE user_id = ? ORDER BY timestamp DESC", (user_id,))
-    return cursor.fetchall()
 
-def db_remove_file(enc_path):
+def db_get_user_files(user_id: int) -> List[Tuple[str, str, int]]:
+    cursor.execute(
+        "SELECT enc_path, orig_name, size FROM files WHERE user_id = ? ORDER BY timestamp DESC",
+        (user_id,),
+    )
+    return [(row[0], row[1], row[2]) for row in cursor.fetchall()]
+
+
+def db_remove_file(enc_path: str) -> None:
     cursor.execute("DELETE FROM files WHERE enc_path = ?", (enc_path,))
     conn.commit()
 
-def db_remove_all_user_files(user_id):
+
+def db_remove_all_user_files(user_id: int) -> None:
     cursor.execute("DELETE FROM files WHERE user_id = ?", (user_id,))
     conn.commit()
 
-# --------------- CALLBACK MAPPINGS ---------------
-file_callback_map = {}     # token -> enc_path
-upload_prompt_map = {}     # token -> list of (user_id, message_id, orig_filename)
 
-# Album group storage: media_group_id -> dict
-album_groups = {}
+# =============================================================================
+# CALLBACK TOKEN MAPS
+# =============================================================================
 
-def make_token():
+file_callback_map: dict[str, str] = {}          # token → enc_path
+# Each entry: tuple with 3 or 4 elements.
+# Local: (user_id, msg_id, filename, file_id)
+# Remote: (user_id, file_id, filename)
+upload_prompt_map: dict[str, List[Tuple]] = {}
+
+# Album grouping
+album_groups: dict[str, dict] = {}
+
+
+def make_token() -> str:
     return secrets.token_hex(8)
 
-# --------------- FLASK HEALTH SERVER ---------------
+
+# =============================================================================
+# FLASK HEALTH SERVER
+# =============================================================================
+
 web = Flask(__name__)
 
+
 @web.route("/")
 def home():
-    return "Telegram HF Storage Bot Running"
+    return "Status: OK - Bot Running"
+
 
 def run_web():
     web.run(host="0.0.0.0", port=7860)
 
-# --------------- HUGGING FACE CLIENT ---------------
+
+# =============================================================================
+# HUGGING FACE CLIENT & PYROGRAM BOT
+# =============================================================================
+
 hf = HfApi(token=HF_TOKEN)
 
-# --------------- PYROGRAM BOT ---------------
 app = Client(
     "storage_bot",
     api_id=API_ID,
     api_hash=API_HASH,
-    bot_token=BOT_TOKEN
+    bot_token=BOT_TOKEN,
 )
 
-# --------------- ASYNC QUEUE & STATE ---------------
-upload_queue = asyncio.Queue()
-active_jobs = {}
+# =============================================================================
+# ASYNC QUEUE & STATE
+# =============================================================================
+
+upload_queue: asyncio.Queue = asyncio.Queue()
+active_jobs: dict[int, dict] = {}
 bot_start_time = time.time()
 
-# --------------- HELPERS ---------------
-def cleanup(path: str):
+
+# =============================================================================
+# UTILITY FUNCTIONS
+# =============================================================================
+
+def cleanup(path: str) -> None:
     try:
         if os.path.isdir(path):
             shutil.rmtree(path, ignore_errors=True)
@@ -132,36 +208,39 @@ def cleanup(path: str):
     except Exception:
         pass
 
+
 def get_free_disk() -> int:
-    st = shutil.disk_usage("/")
-    return st.free
+    return shutil.disk_usage("/").free
+
 
-def format_size(bytes_val):
+def format_size(bytes_val: int) -> str:
     for unit in ['B', 'KB', 'MB', 'GB', 'TB']:
         if bytes_val < 1024.0:
             return f"{bytes_val:.1f} {unit}"
         bytes_val /= 1024.0
     return f"{bytes_val:.1f} PB"
 
-def format_speed(bytes_per_sec):
+
+def format_speed(bytes_per_sec: float) -> str:
     return f"{format_size(bytes_per_sec)}/s"
 
+
 def get_memory_usage() -> str:
-    """Return current process memory usage (RSS) as a human-readable string."""
     try:
-        # Linux only
         with open("/proc/self/status") as f:
             for line in f:
                 if line.startswith("VmRSS:"):
-                    # Value is in kB
-                    kb = int(line.split()[1])
-                    return format_size(kb * 1024)
+                    return format_size(int(line.split()[1]) * 1024)
     except Exception:
-        pass
-    return "unknown"
+        return "unknown"
+
 
-async def send_progress_updater(job_id, status_msg, cancel_event, total_bytes):
-    """Shows Telegram send progress every 3s."""
+async def send_progress_updater(
+    job_id: int,
+    status_msg: Message,
+    cancel_event: asyncio.Event,
+    total_bytes: int,
+) -> Tuple[asyncio.Task, list]:
     bytes_holder = [0]
     start_time = time.time()
 
@@ -171,13 +250,20 @@ async def send_progress_updater(job_id, status_msg, cancel_event, total_bytes):
             current = bytes_holder[0]
             speed = current / elapsed
             pct = (100 * current / total_bytes) if total_bytes else 0
-            text = f"📤 Sending to Telegram...\n{pct:.1f}% of {format_size(total_bytes)}\n{format_size(current)} sent\nSpeed: {format_speed(speed)}\nETA: {int((total_bytes - current) / speed) if speed > 0 else '...'}s"
+            eta = int((total_bytes - current) / speed) if speed > 0 else -1
+            text = (
+                f"📤 Sending to Telegram...\n"
+                f"{pct:.1f}% of {format_size(total_bytes)}\n"
+                f"{format_size(current)} sent\n"
+                f"Speed: {format_speed(speed)}\n"
+                f"ETA: {eta}s" if eta >= 0 else "ETA: ..."
+            )
             try:
                 await status_msg.edit_text(
                     text,
-                    reply_markup=InlineKeyboardMarkup([
-                        [InlineKeyboardButton("❌ Cancel", callback_data=f"cancel_job:{job_id}")]
-                    ])
+                    reply_markup=InlineKeyboardMarkup(
+                        [[InlineKeyboardButton("❌ Cancel", callback_data=f"cancel_job:{job_id}")]]
+                    ),
                 )
             except Exception:
                 break
@@ -186,16 +272,31 @@ async def send_progress_updater(job_id, status_msg, cancel_event, total_bytes):
     task = asyncio.create_task(updater())
     return task, bytes_holder
 
-# --------------- UPLOAD JOB ---------------
-async def do_upload_job(user_id, chat_id, file_msg, orig_filename, job_id, cancel_event, status_msg):
+
+# =============================================================================
+# UPLOAD / DOWNLOAD JOBS
+# =============================================================================
+
+async def do_upload_job(
+    user_id: int,
+    chat_id: int,
+    file_msg_or_id: Union[Message, str],
+    orig_filename: str,
+    job_id: int,
+    cancel_event: asyncio.Event,
+    status_msg: Message,
+    file_id: Optional[str] = None,
+) -> None:
+    """Download file from Telegram (by message or file_id) and upload to HF."""
     temp_dir = tempfile.mkdtemp()
     local_path = None
     try:
-        # 1. Download from Telegram (show speed/progress)
-        await status_msg.edit_text("⬇️ Downloading from Telegram...",
-                                   reply_markup=InlineKeyboardMarkup([
-                                       [InlineKeyboardButton("❌ Cancel", callback_data=f"cancel_job:{job_id}")]
-                                   ]))
+        await status_msg.edit_text(
+            "⬇️ Downloading from Telegram...",
+            reply_markup=InlineKeyboardMarkup(
+                [[InlineKeyboardButton("❌ Cancel", callback_data=f"cancel_job:{job_id}")]]
+            ),
+        )
         last_update = [0]
 
         async def download_progress(current, total):
@@ -206,36 +307,52 @@ async def do_upload_job(user_id, chat_id, file_msg, orig_filename, job_id, cance
             elapsed = max(time.time() - download_start, 0.01)
             speed = current / elapsed
             pct = (100 * current / total) if total else 0
-            text = f"⬇️ Downloading from Telegram...\n{pct:.1f}% of {format_size(total)}\n{format_size(current)} downloaded\nSpeed: {format_speed(speed)}"
+            text = (
+                f"⬇️ Downloading from Telegram...\n"
+                f"{pct:.1f}% of {format_size(total)}\n"
+                f"{format_size(current)} downloaded\n"
+                f"Speed: {format_speed(speed)}"
+            )
             try:
                 await status_msg.edit_text(
                     text,
-                    reply_markup=InlineKeyboardMarkup([
-                        [InlineKeyboardButton("❌ Cancel", callback_data=f"cancel_job:{job_id}")]
-                    ])
+                    reply_markup=InlineKeyboardMarkup(
+                        [[InlineKeyboardButton("❌ Cancel", callback_data=f"cancel_job:{job_id}")]]
+                    ),
                 )
             except Exception:
                 pass
 
         download_start = time.time()
-        local_path = await file_msg.download(
-            file_name=os.path.join(temp_dir, orig_filename),
-            progress=download_progress
-        )
+
+        if isinstance(file_msg_or_id, Message):
+            local_path = await file_msg_or_id.download(
+                file_name=os.path.join(temp_dir, orig_filename),
+                progress=download_progress,
+            )
+        else:
+            # Download by file_id
+            local_path = await app.download_media(
+                file_msg_or_id,
+                file_name=os.path.join(temp_dir, orig_filename),
+                progress=download_progress,
+            )
+
         if cancel_event.is_set():
             return
 
         file_size = os.path.getsize(local_path)
 
-        # 2. Upload to Hugging Face (no speed/progress)
         enc_folder = get_encrypted_folder(user_id)
         enc_filename = get_encrypted_filename(orig_filename)
         repo_path = f"{enc_folder}/{enc_filename}"
 
-        await status_msg.edit_text("⬆️ Uploading to Hugging Face...",
-                                   reply_markup=InlineKeyboardMarkup([
-                                       [InlineKeyboardButton("❌ Cancel", callback_data=f"cancel_job:{job_id}")]
-                                   ]))
+        await status_msg.edit_text(
+            "⬆️ Uploading to Hugging Face...",
+            reply_markup=InlineKeyboardMarkup(
+                [[InlineKeyboardButton("❌ Cancel", callback_data=f"cancel_job:{job_id}")]]
+            ),
+        )
 
         loop = asyncio.get_event_loop()
         upload_done = loop.create_future()
@@ -245,7 +362,7 @@ async def do_upload_job(user_id, chat_id, file_msg, orig_filename, job_id, cance
                 path_or_fileobj=local_path,
                 path_in_repo=repo_path,
                 repo_id=DATASET_REPO,
-                repo_type="dataset",
+                repo_type=DATASET_TYPE,
             )
 
         task = loop.run_in_executor(None, sync_upload)
@@ -260,7 +377,7 @@ async def do_upload_job(user_id, chat_id, file_msg, orig_filename, job_id, cance
                 hf.delete_file(
                     path_in_repo=repo_path,
                     repo_id=DATASET_REPO,
-                    repo_type="dataset",
+                    repo_type=DATASET_TYPE,
                 )
             except Exception:
                 pass
@@ -270,33 +387,43 @@ async def do_upload_job(user_id, chat_id, file_msg, orig_filename, job_id, cance
         await task
         db_add_file(user_id, repo_path, orig_filename, file_size)
         await status_msg.edit_text(f"✅ Uploaded: `{orig_filename}`", reply_markup=None)
+        logger.info("Uploaded %s for user %d", orig_filename, user_id)
 
     except Exception as e:
+        logger.exception("Upload failed for %s", orig_filename)
         await status_msg.edit_text(f"❌ Upload failed: {e}", reply_markup=None)
     finally:
         cleanup(temp_dir)
 
-# --------------- DOWNLOAD JOB ---------------
-async def do_download_job(chat_id, enc_path, orig_name, job_id, cancel_event, status_msg):
+
+async def do_download_job(
+    chat_id: int,
+    enc_path: str,
+    orig_name: str,
+    job_id: int,
+    cancel_event: asyncio.Event,
+    status_msg: Message,
+) -> None:
     temp_dir = tempfile.mkdtemp()
     try:
         loop = asyncio.get_event_loop()
 
-        # 1. Download from HF using robust hf_hub_download (resilient to connection issues)
-        await status_msg.edit_text("⬇️ Downloading from dataset...",
-                                   reply_markup=InlineKeyboardMarkup([
-                                       [InlineKeyboardButton("❌ Cancel", callback_data=f"cancel_job:{job_id}")]
-                                   ]))
+        await status_msg.edit_text(
+            "⬇️ Downloading from dataset...",
+            reply_markup=InlineKeyboardMarkup(
+                [[InlineKeyboardButton("❌ Cancel", callback_data=f"cancel_job:{job_id}")]]
+            ),
+        )
 
         download_future = loop.run_in_executor(
             None,
             lambda: hf_hub_download(
                 repo_id=DATASET_REPO,
                 filename=enc_path,
-                repo_type="dataset",
+                repo_type=DATASET_TYPE,
                 token=HF_TOKEN,
                 local_dir=temp_dir,
-            )
+            ),
         )
 
         while not download_future.done():
@@ -308,7 +435,6 @@ async def do_download_job(chat_id, enc_path, orig_name, job_id, cancel_event, st
         local_file = download_future.result()
         total_size = os.path.getsize(local_file)
 
-        # 2. Send back to Telegram (with speed/progress if single file)
         if total_size <= MAX_FILE_SIZE_TG:
             prog_task, bytes_holder = await send_progress_updater(
                 job_id, status_msg, cancel_event, total_size
@@ -322,7 +448,7 @@ async def do_download_job(chat_id, enc_path, orig_name, job_id, cancel_event, st
                 document=local_file,
                 file_name=orig_name,
                 caption=orig_name,
-                progress=send_progress
+                progress=send_progress,
             )
             prog_task.cancel()
             await status_msg.delete()
@@ -334,32 +460,44 @@ async def do_download_job(chat_id, enc_path, orig_name, job_id, cancel_event, st
                     if cancel_event.is_set():
                         await status_msg.edit_text("❌ Cancelled during send.")
                         return
-                    await status_msg.edit_text(f"📤 Sending part {i}/{num_parts}...",
-                                               reply_markup=InlineKeyboardMarkup([
-                                                   [InlineKeyboardButton("❌ Cancel", callback_data=f"cancel_job:{job_id}")]
-                                               ]))
+                    await status_msg.edit_text(
+                        f"📤 Sending part {i}/{num_parts}...",
+                        reply_markup=InlineKeyboardMarkup(
+                            [[InlineKeyboardButton("❌ Cancel", callback_data=f"cancel_job:{job_id}")]]
+                        ),
+                    )
                     chunk = f.read(part_size)
                     part_name = f"{orig_name}.part{i}of{num_parts}"
                     part_path = os.path.join(temp_dir, part_name)
                     with open(part_path, "wb") as p:
                         p.write(chunk)
-                    await app.send_document(chat_id, document=part_path, file_name=part_name,
-                                           caption=f"{orig_name} – Part {i}/{num_parts}")
+                    await app.send_document(
+                        chat_id,
+                        document=part_path,
+                        file_name=part_name,
+                        caption=f"{orig_name} – Part {i}/{num_parts}",
+                    )
                     os.remove(part_path)
             await status_msg.delete()
+        logger.info("Downloaded %s for chat %d", orig_name, chat_id)
 
     except Exception as e:
+        logger.exception("Download failed for %s", orig_name)
         await status_msg.edit_text(f"❌ Download failed: {e}")
     finally:
         cleanup(temp_dir)
 
-# --------------- DELETE FUNCTIONS ---------------
-async def delete_user_file(chat_id, user_id, enc_path, token, query):
+
+# =============================================================================
+# DELETE FUNCTIONS
+# =============================================================================
+
+async def delete_user_file(chat_id, user_id, enc_path, token, query: CallbackQuery):
     try:
         hf.delete_file(
             path_in_repo=enc_path,
             repo_id=DATASET_REPO,
-            repo_type="dataset",
+            repo_type=DATASET_TYPE,
         )
         db_remove_file(enc_path)
         file_callback_map.pop(token, None)
@@ -368,7 +506,8 @@ async def delete_user_file(chat_id, user_id, enc_path, token, query):
     except Exception as e:
         await query.answer(f"Error: {e}", show_alert=True)
 
-async def delete_all_user_files(chat_id, user_id, query):
+
+async def delete_all_user_files(chat_id, user_id, query: CallbackQuery):
     user_files = db_get_user_files(user_id)
     if not user_files:
         await query.answer("No files to delete.", show_alert=True)
@@ -379,42 +518,50 @@ async def delete_all_user_files(chat_id, user_id, query):
             hf.delete_file(
                 path_in_repo=enc_path,
                 repo_id=DATASET_REPO,
-                repo_type="dataset",
+                repo_type=DATASET_TYPE,
             )
         except Exception:
             pass
 
     db_remove_all_user_files(user_id)
 
-    tokens_to_delete = [tok for tok, path in file_callback_map.items() if path.startswith(get_encrypted_folder(user_id))]
+    tokens_to_delete = [
+        tok for tok, path in file_callback_map.items()
+        if path.startswith(get_encrypted_folder(user_id))
+    ]
     for tok in tokens_to_delete:
         file_callback_map.pop(tok, None)
 
     await query.answer(f"Deleted {len(user_files)} file(s).", show_alert=True)
     await refresh_file_list(chat_id, user_id, query.message)
 
-async def refresh_file_list(chat_id, user_id, message):
+
+async def refresh_file_list(chat_id: int, user_id: int, message: Message):
     user_files = db_get_user_files(user_id)
     if not user_files:
         return await message.edit_text("You have no uploaded files.", reply_markup=None)
 
     buttons = []
-    for enc_path, orig_name, size in user_files:
+    for enc_path, orig_name, _ in user_files:
         token = make_token()
         file_callback_map[token] = enc_path
         buttons.append([
             InlineKeyboardButton(f"📥 {orig_name[:40]}", callback_data=f"get:{token}"),
-            InlineKeyboardButton("🗑 Delete", callback_data=f"delete:{token}")
+            InlineKeyboardButton("🗑 Delete", callback_data=f"delete:{token}"),
         ])
 
     buttons.append([InlineKeyboardButton("⚠️ Delete All Files", callback_data=f"delete_all:{user_id}")])
 
     await message.edit_text(
         "Your uploaded files (click to download, delete to remove):",
-        reply_markup=InlineKeyboardMarkup(buttons)
+        reply_markup=InlineKeyboardMarkup(buttons),
     )
 
-# --------------- ALBUM GROUP PROCESSING ---------------
+
+# =============================================================================
+# ALBUM GROUPING
+# =============================================================================
+
 async def process_album_group(media_group_id: str):
     group_data = album_groups.pop(media_group_id, None)
     if not group_data:
@@ -422,97 +569,144 @@ async def process_album_group(media_group_id: str):
 
     chat_id = group_data['chat_id']
     user_id = group_data['user_id']
-    messages = group_data['messages']   # list of (message_id, filename)
+    messages = group_data['messages']  # list of (message_id, filename, file_id)
 
     if len(messages) == 1:
-        msg_id, filename = messages[0]
+        msg_id, filename, file_id = messages[0]
         token = make_token()
-        upload_prompt_map[token] = [(user_id, msg_id, filename)]
+        upload_prompt_map[token] = [(user_id, msg_id, filename, file_id)]
         kb = InlineKeyboardMarkup([
             [InlineKeyboardButton("⬆️ Upload File", callback_data=f"upload:{token}")],
-            [InlineKeyboardButton("❌ Cancel", callback_data="cancel_upload:0")]
+            [InlineKeyboardButton("❌ Cancel", callback_data="cancel_upload:0")],
         ])
         await app.send_message(chat_id, f"Do you want to upload `{filename}`?", reply_markup=kb)
     else:
-        files_list = "\n".join(f"• {fname}" for _, fname in messages)
+        files_list = "\n".join(f"• {fname}" for _, fname, _ in messages)
         token = make_token()
-        upload_prompt_map[token] = [(user_id, msg_id, fname) for msg_id, fname in messages]
+        upload_prompt_map[token] = [(user_id, msg_id, fname, fid) for msg_id, fname, fid in messages]
 
-        button_rows = []
-        button_rows.append([InlineKeyboardButton("⬆️ Upload All", callback_data=f"upload_all:{token}")])
-        for msg_id, fname in messages:
+        button_rows = [
+            [InlineKeyboardButton("⬆️ Upload All", callback_data=f"upload_all:{token}")]
+        ]
+        for msg_id, fname, fid in messages:
             individual_token = make_token()
-            upload_prompt_map[individual_token] = [(user_id, msg_id, fname)]
-            button_rows.append([InlineKeyboardButton(f"📤 {fname[:40]}", callback_data=f"upload:{individual_token}")])
+            upload_prompt_map[individual_token] = [(user_id, msg_id, fname, fid)]
+            button_rows.append(
+                [InlineKeyboardButton(f"📤 {fname[:40]}", callback_data=f"upload:{individual_token}")]
+            )
         button_rows.append([InlineKeyboardButton("❌ Cancel", callback_data="cancel_upload:0")])
 
         await app.send_message(
             chat_id,
             f"Received {len(messages)} files.\n{files_list}\n\nChoose:",
-            reply_markup=InlineKeyboardMarkup(button_rows)
+            reply_markup=InlineKeyboardMarkup(button_rows),
         )
 
-# --------------- QUEUE WORKER ---------------
+
+# =============================================================================
+# UPLOAD QUEUE WORKER
+# =============================================================================
+
 async def upload_worker():
     while True:
         job = await upload_queue.get()
-        user_id, chat_id, file_msg, orig_filename, job_id = job
+        user_id, chat_id, file_obj, orig_filename, job_id, file_id = job
         cancel_event = asyncio.Event()
         active_jobs[job_id] = {"status": "running", "cancel_event": cancel_event}
 
         while get_free_disk() < MIN_FREE_DISK:
+            logger.warning("Low disk space – pausing uploads")
             await asyncio.sleep(10)
 
         status_msg = await app.send_message(
-            chat_id, "⏳ Starting upload...",
-            reply_markup=InlineKeyboardMarkup([
-                [InlineKeyboardButton("❌ Cancel", callback_data=f"cancel_job:{job_id}")]
-            ])
+            chat_id,
+            "⏳ Starting upload...",
+            reply_markup=InlineKeyboardMarkup(
+                [[InlineKeyboardButton("❌ Cancel", callback_data=f"cancel_job:{job_id}")]]
+            ),
         )
         active_jobs[job_id]["status_msg"] = status_msg
 
         try:
-            await do_upload_job(user_id, chat_id, file_msg, orig_filename, job_id, cancel_event, status_msg)
+            await do_upload_job(
+                user_id, chat_id, file_obj, orig_filename, job_id, cancel_event, status_msg, file_id
+            )
         finally:
             active_jobs.pop(job_id, None)
             upload_queue.task_done()
 
-# --------------- BOT HANDLERS ---------------
+
+# =============================================================================
+# BOT COMMAND HANDLERS
+# =============================================================================
+
 @app.on_message(filters.command("start"))
-async def start_cmd(_, message):
+async def cmd_start(_, message: Message):
     await message.reply_text(
-        "Send me any file or a Telegram message link to upload.\n"
-        "/files – list your uploaded files\n"
-        "/status – show disk, dataset & system status"
+        "👋 **Welcome to the Storage Bot!**\n\n"
+        "Send me any file or a Telegram message link to upload it to our cloud dataset.\n"
+        "Use /help for a full command list."
     )
 
+
+@app.on_message(filters.command("help"))
+async def cmd_help(_, message: Message):
+    text = (
+        "**Available Commands:**\n\n"
+        "/start – Show welcome message\n"
+        "/files – List your uploaded files (with download/delete buttons)\n"
+        "/status – Show system & dataset space\n"
+        "/delete_all – Delete all your files (confirmation button)\n"
+        "/upload <link> – Directly upload a file from a Telegram message link\n"
+        "/help – This help text\n\n"
+        "**How to use:**\n"
+        "• Send any file (document, video, audio, photo, sticker, voice, etc.) → bot asks to upload.\n"
+        "• Send multiple files as an album → bot shows one message with 'Upload All' + individual buttons.\n"
+        "• Send a message link (like https://t.me/.../123) → bot fetches the file and asks to upload.\n"
+        "• Use /upload <link> to skip the prompt and upload instantly.\n"
+        "• In /files list, click a file to download it back, or use the Delete button."
+    )
+    await message.reply_text(text, disable_web_page_preview=True)
+
+
 @app.on_message(filters.command("status"))
-async def status_cmd(_, message):
+async def cmd_status(_, message: Message):
     free = get_free_disk()
     disk_text = f"💾 Local disk: {format_size(free)} free (min {format_size(MIN_FREE_DISK)})"
 
     try:
         info = hf.dataset_info(repo_id=DATASET_REPO)
-        used = info.used_storage if info.used_storage else 0
+        used = info.used_storage or 0
     except Exception:
         used = 0
     dataset_text = f"📦 Dataset: {format_size(used)} used of {format_size(DATASET_MAX_SIZE)}"
 
-    uptime_seconds = int(time.time() - bot_start_time)
-    uptime_text = f"⏱️ Uptime: {uptime_seconds//3600}h {(uptime_seconds%3600)//60}m {uptime_seconds%60}s"
+    uptime_sec = int(time.time() - bot_start_time)
+    uptime_text = f"⏱️ Uptime: {uptime_sec//3600}h {(uptime_sec%3600)//60}m {uptime_sec%60}s"
     mem_text = f"🧠 Memory: {get_memory_usage()}"
+    cpu_text = f"🖥️ CPU cores: {os.cpu_count()}"
+
+    try:
+        r = requests.head("https://huggingface.co", timeout=5)
+        hf_ping = "OK" if r.status_code < 400 else f"HTTP {r.status_code}"
+    except Exception:
+        hf_ping = "unreachable"
+
+    await message.reply_text(
+        f"{disk_text}\n{dataset_text}\n{uptime_text}\n{mem_text}\n{cpu_text}\n"
+        f"🌐 Hugging Face: {hf_ping}"
+    )
 
-    await message.reply_text(f"{disk_text}\n{dataset_text}\n{uptime_text}\n{mem_text}")
 
 @app.on_message(filters.command("files"))
-async def list_files_cmd(_, message):
+async def cmd_files(_, message: Message):
     user_id = message.from_user.id
     user_files = db_get_user_files(user_id)
 
     if not user_files:
         enc_folder = get_encrypted_folder(user_id)
         try:
-            all_hf = hf.list_repo_files(repo_id=DATASET_REPO, repo_type="dataset")
+            all_hf = hf.list_repo_files(repo_id=DATASET_REPO, repo_type=DATASET_TYPE)
             hf_user_files = [f for f in all_hf if f.startswith(enc_folder + "/")]
             for path in hf_user_files:
                 _, enc_name = path.split("/", 1)
@@ -521,41 +715,123 @@ async def list_files_cmd(_, message):
                 except Exception:
                     orig_name = enc_name
                 try:
-                    file_url = hf_hub_url(repo_id=DATASET_REPO, filename=path, repo_type="dataset")
+                    file_url = hf_hub_url(repo_id=DATASET_REPO, filename=path, repo_type=DATASET_TYPE)
                     meta = get_hf_file_metadata(url=file_url, token=HF_TOKEN)
                     size = meta.size
                 except Exception:
                     size = 0
                 db_add_file(user_id, path, orig_name, size)
             user_files = db_get_user_files(user_id)
-        except Exception:
-            pass
+        except Exception as e:
+            logger.warning("Fallback scan failed for user %d: %s", user_id, e)
 
     if not user_files:
         return await message.reply_text("You have no uploaded files.")
 
     buttons = []
-    for enc_path, orig_name, size in user_files:
+    for enc_path, orig_name, _ in user_files:
         token = make_token()
         file_callback_map[token] = enc_path
         buttons.append([
             InlineKeyboardButton(f"📥 {orig_name[:40]}", callback_data=f"get:{token}"),
-            InlineKeyboardButton("🗑 Delete", callback_data=f"delete:{token}")
+            InlineKeyboardButton("🗑 Delete", callback_data=f"delete:{token}"),
         ])
 
     buttons.append([InlineKeyboardButton("⚠️ Delete All Files", callback_data=f"delete_all:{user_id}")])
 
     await message.reply_text(
         "Your uploaded files (click to download, delete to remove):",
-        reply_markup=InlineKeyboardMarkup(buttons)
+        reply_markup=InlineKeyboardMarkup(buttons),
     )
 
+
+@app.on_message(filters.command("delete_all"))
+async def cmd_delete_all(_, message: Message):
+    user_id = message.from_user.id
+    kb = InlineKeyboardMarkup([
+        [InlineKeyboardButton("Yes, delete all my files", callback_data=f"delete_all:{user_id}")],
+        [InlineKeyboardButton("Cancel", callback_data="cancel_upload:0")],
+    ])
+    await message.reply_text(
+        "⚠️ This will permanently delete **all** your files from the dataset. Are you sure?",
+        reply_markup=kb,
+    )
+
+
+@app.on_message(filters.command("upload"))
+async def cmd_upload(_, message: Message):
+    args = message.text.split(maxsplit=1)
+    if len(args) < 2:
+        return await message.reply_text("Usage: /upload <link>")
+
+    link = args[1].strip()
+    match = re.search(r'(https?://t\.me/((?:c/)?[\w\d_]+)/(\d+))', link)
+    if not match:
+        return await message.reply_text("Invalid Telegram message link.")
+
+    chat_str, msg_id = match.group(2), int(match.group(3))
+    try:
+        if chat_str.startswith("c/"):
+            chat_id = int("-100" + chat_str[2:])
+        else:
+            chat_id = chat_str
+        msg = await app.get_messages(chat_id, msg_id)
+    except Exception as e:
+        return await message.reply_text(f"❌ Could not fetch the message: {e}")
+
+    if not msg.media:
+        return await message.reply_text("The linked message doesn't contain any downloadable file.")
+
+    # Extract file_id and filename
+    file_id = None
+    if msg.document:
+        filename = msg.document.file_name
+        file_id = msg.document.file_id
+    elif msg.video:
+        filename = msg.video.file_name or f"{msg.video.file_unique_id}.mp4"
+        file_id = msg.video.file_id
+    elif msg.audio:
+        filename = msg.audio.file_name or f"{msg.audio.file_unique_id}.mp3"
+        file_id = msg.audio.file_id
+    elif msg.photo:
+        filename = f"{msg.photo.file_unique_id}.jpg"
+        file_id = msg.photo.file_id
+    elif msg.animation:
+        filename = msg.animation.file_name or f"{msg.animation.file_unique_id}.mp4"
+        file_id = msg.animation.file_id
+    elif msg.sticker:
+        if msg.sticker.is_animated:
+            filename = f"{msg.sticker.file_unique_id}.tgs"
+        elif msg.sticker.is_video:
+            filename = f"{msg.sticker.file_unique_id}.webm"
+        else:
+            filename = f"{msg.sticker.file_unique_id}.webp"
+        file_id = msg.sticker.file_id
+    elif msg.video_note:
+        filename = f"{msg.video_note.file_unique_id}.mp4"
+        file_id = msg.video_note.file_id
+    elif msg.voice:
+        filename = f"{msg.voice.file_unique_id}.ogg"
+        file_id = msg.voice.file_id
+    else:
+        filename = f"file_{msg_id}"
+
+    job_id = int(time.time() * 1000 + hash(filename))
+    active_jobs[job_id] = {"status": "queued"}
+    await upload_queue.put((message.from_user.id, message.chat.id, file_id, filename, job_id, file_id))
+    await message.reply_text(f"📦 `{filename}` queued for upload.")
+
+
+# =============================================================================
+# CALLBACK HANDLERS
+# =============================================================================
+
 @app.on_callback_query(filters.regex("^get:"))
-async def get_file_callback(_, query):
+async def cb_get_file(_, query: CallbackQuery):
     token = query.data.replace("get:", "")
     enc_path = file_callback_map.get(token)
     if not enc_path:
-        return await query.answer("Invalid or expired file link. Use /files again.", show_alert=True)
+        return await query.answer("Invalid or expired link. Use /files again.", show_alert=True)
 
     job_id = int(time.time() * 1000)
     cancel_event = asyncio.Event()
@@ -568,41 +844,33 @@ async def get_file_callback(_, query):
 
     status_msg = await query.message.edit_text(
         "⏳ Preparing download...",
-        reply_markup=InlineKeyboardMarkup([
-            [InlineKeyboardButton("❌ Cancel", callback_data=f"cancel_job:{job_id}")]
-        ])
+        reply_markup=InlineKeyboardMarkup(
+            [[InlineKeyboardButton("❌ Cancel", callback_data=f"cancel_job:{job_id}")]]
+        ),
     )
     active_jobs[job_id]["status_msg"] = status_msg
 
-    asyncio.create_task(do_download_job(
-        query.message.chat.id,
-        enc_path,
-        orig_name,
-        job_id,
-        cancel_event,
-        status_msg
-    ))
+    asyncio.create_task(do_download_job(query.message.chat.id, enc_path, orig_name, job_id, cancel_event, status_msg))
+
 
 @app.on_callback_query(filters.regex("^delete:"))
-async def delete_file_callback(_, query):
+async def cb_delete_file(_, query: CallbackQuery):
     token = query.data.replace("delete:", "")
     enc_path = file_callback_map.pop(token, None)
     if not enc_path:
         return await query.answer("File not found or already deleted.", show_alert=True)
-
     user_id = query.from_user.id
-    chat_id = query.message.chat.id
+    await delete_user_file(query.message.chat.id, user_id, enc_path, token, query)
 
-    await delete_user_file(chat_id, user_id, enc_path, token, query)
 
 @app.on_callback_query(filters.regex("^delete_all:"))
-async def delete_all_callback(_, query):
+async def cb_delete_all(_, query: CallbackQuery):
     user_id = int(query.data.split(":")[-1])
-    chat_id = query.message.chat.id
-    await delete_all_user_files(chat_id, user_id, query)
+    await delete_all_user_files(query.message.chat.id, user_id, query)
+
 
 @app.on_callback_query(filters.regex("^cancel_job:"))
-async def cancel_job_callback(_, query):
+async def cb_cancel_job(_, query: CallbackQuery):
     _, job_id_str = query.data.split(":")
     job_id = int(job_id_str)
     if job_id in active_jobs:
@@ -611,81 +879,100 @@ async def cancel_job_callback(_, query):
     else:
         await query.answer("Job not found.")
 
+
 @app.on_callback_query(filters.regex("^upload:"))
-async def upload_callback(_, query):
+async def cb_upload(_, query: CallbackQuery):
     token = query.data.replace("upload:", "")
-    data = upload_prompt_map.pop(token, None)
-    if not data:
-        return await query.answer("Upload prompt expired. Please send the file again.", show_alert=True)
+    entries = upload_prompt_map.pop(token, None)
+    if not entries:
+        return await query.answer("Prompt expired.", show_alert=True)
 
-    chat_id = query.message.chat.id
-    for user_id, msg_id, orig_filename in data:
-        try:
-            file_msg = await app.get_messages(chat_id, msg_id)
-        except Exception:
-            await query.answer(f"Message for {orig_filename} missing.", show_alert=True)
-            continue
-        if not file_msg:
-            await query.answer(f"File {orig_filename} no longer available.", show_alert=True)
-            continue
+    for entry in entries:
+        user_id = entry[0]
+        if len(entry) == 4:
+            # Local: (user_id, msg_id, filename, file_id)
+            msg_id = entry[1]
+            filename = entry[2]
+            file_id = entry[3]
+            try:
+                file_msg = await app.get_messages(query.message.chat.id, msg_id)
+            except Exception:
+                await query.answer(f"Message for {filename} missing.", show_alert=True)
+                continue
+            if not file_msg:
+                await query.answer(f"File {filename} no longer available.", show_alert=True)
+                continue
+            file_obj = file_msg
+        else:
+            # Remote: (user_id, file_id, filename)
+            file_id = entry[1]
+            filename = entry[2]
+            file_obj = file_id
 
-        job_id = int(time.time() * 1000 + hash(orig_filename))
+        job_id = int(time.time() * 1000 + hash(filename))
         active_jobs[job_id] = {"status": "queued"}
-        await upload_queue.put((user_id, chat_id, file_msg, orig_filename, job_id))
+        await upload_queue.put((user_id, query.message.chat.id, file_obj, filename, job_id, file_id))
+
+    await query.message.edit_text(f"📦 {len(entries)} file(s) added to queue.", reply_markup=None)
 
-    await query.message.edit_text(
-        f"📦 {len(data)} file(s) added to queue.",
-        reply_markup=None
-    )
 
 @app.on_callback_query(filters.regex("^upload_all:"))
-async def upload_all_callback(_, query):
+async def cb_upload_all(_, query: CallbackQuery):
     token = query.data.replace("upload_all:", "")
-    data = upload_prompt_map.pop(token, None)
-    if not data:
+    entries = upload_prompt_map.pop(token, None)
+    if not entries:
         return await query.answer("Prompt expired.", show_alert=True)
 
-    chat_id = query.message.chat.id
-    for user_id, msg_id, orig_filename in data:
-        try:
-            file_msg = await app.get_messages(chat_id, msg_id)
-        except Exception:
-            continue
-        if not file_msg:
-            continue
-        job_id = int(time.time() * 1000 + hash(orig_filename))
+    for entry in entries:
+        user_id = entry[0]
+        if len(entry) == 4:
+            msg_id = entry[1]
+            filename = entry[2]
+            file_id = entry[3]
+            try:
+                file_msg = await app.get_messages(query.message.chat.id, msg_id)
+            except Exception:
+                continue
+            if not file_msg:
+                continue
+            file_obj = file_msg
+        else:
+            file_id = entry[1]
+            filename = entry[2]
+            file_obj = file_id
+
+        job_id = int(time.time() * 1000 + hash(filename))
         active_jobs[job_id] = {"status": "queued"}
-        await upload_queue.put((user_id, chat_id, file_msg, orig_filename, job_id))
+        await upload_queue.put((user_id, query.message.chat.id, file_obj, filename, job_id, file_id))
+
+    await query.message.edit_text(f"📦 All {len(entries)} files added to queue.", reply_markup=None)
 
-    await query.message.edit_text(
-        f"📦 All {len(data)} files added to queue.",
-        reply_markup=None
-    )
 
 @app.on_callback_query(filters.regex("^cancel_upload:"))
-async def cancel_upload_prompt(_, query):
+async def cb_cancel_upload_prompt(_, query: CallbackQuery):
     await query.message.delete()
 
-# ----------- MESSAGE LINK HANDLER -----------
-@app.on_message(filters.text & ~filters.command(["start", "files", "status", "delete_all"]) & filters.regex(r'https?://t\.me/'))
-async def on_message_link(client, message):
-    # Extract first valid t.me link
+
+# =============================================================================
+# MESSAGE LINK HANDLER (interactive prompt)
+# =============================================================================
+
+@app.on_message(
+    filters.text
+    & ~filters.command(["start", "help", "status", "files", "delete_all", "upload"])
+    & filters.regex(r'https?://t\.me/')
+)
+async def on_message_link(client, message: Message):
     links = re.findall(r'(https?://t\.me/((?:c/)?[\w\d_]+)/(\d+))', message.text)
     if not links:
-        return  # shouldn't happen because of filter, but just in case
-
-    # Process the first link found
-    full_link, chat_str, msg_id_str = links[0]
-    msg_id = int(msg_id_str)
+        return
 
+    chat_str, msg_id = links[0][1], int(links[0][2])
     try:
         if chat_str.startswith("c/"):
-            # Numeric chat id: t.me/c/123456789/1234
             chat_id = int("-100" + chat_str[2:])
         else:
-            chat_id = chat_str  # username
-
-        # Fetch the message
+            chat_id = chat_str
         msg = await app.get_messages(chat_id, msg_id)
     except Exception as e:
         await message.reply_text(f"❌ Could not access the message: {e}")
@@ -695,17 +982,23 @@ async def on_message_link(client, message):
         await message.reply_text("The linked message does not contain any downloadable file.")
         return
 
-    # Determine the file name from the fetched message
+    # Extract filename and file_id
+    file_id = None
     if msg.document:
         filename = msg.document.file_name
+        file_id = msg.document.file_id
     elif msg.video:
         filename = msg.video.file_name or f"{msg.video.file_unique_id}.mp4"
+        file_id = msg.video.file_id
     elif msg.audio:
         filename = msg.audio.file_name or f"{msg.audio.file_unique_id}.mp3"
+        file_id = msg.audio.file_id
     elif msg.photo:
         filename = f"{msg.photo.file_unique_id}.jpg"
+        file_id = msg.photo.file_id
     elif msg.animation:
         filename = msg.animation.file_name or f"{msg.animation.file_unique_id}.mp4"
+        file_id = msg.animation.file_id
     elif msg.sticker:
         if msg.sticker.is_animated:
             filename = f"{msg.sticker.file_unique_id}.tgs"
@@ -713,97 +1006,54 @@ async def on_message_link(client, message):
             filename = f"{msg.sticker.file_unique_id}.webm"
         else:
             filename = f"{msg.sticker.file_unique_id}.webp"
+        file_id = msg.sticker.file_id
     elif msg.video_note:
         filename = f"{msg.video_note.file_unique_id}.mp4"
+        file_id = msg.video_note.file_id
     elif msg.voice:
         filename = f"{msg.voice.file_unique_id}.ogg"
+        file_id = msg.voice.file_id
     else:
         filename = f"file_{msg_id}"
 
-    # We'll use the message object directly as the source for download.
-    # Important: The upload queue expects a Telegram message object from the user's chat,
-    # but here the message is from another chat. It still works because Pyrogram's download()
-    # works on any message object the bot can access. So we can pass it directly.
-
+    # Remote entry: (user_id, file_id, filename)
     token = make_token()
-    # Store (user_id, message_id, filename) - but message_id is from the original chat.
-    # For the upload callback we need to retrieve the message by its original chat.
-    # To simplify, we can store the chat_id and msg_id as well.
-    # Let's modify the mapping to include chat_id: (user_id, chat_id_of_file, msg_id, filename)
-    # But the current upload_prompt_map structure expects messages from the same chat as the user.
-    # We'll extend the mapping to handle remote messages. We'll store:
-    # token -> (user_id, file_chat_id, msg_id, filename)
-    # We'll adjust upload_callback accordingly.
-    upload_prompt_map[token] = [(message.from_user.id, chat_id, msg_id, filename)]
+    upload_prompt_map[token] = [(message.from_user.id, file_id, filename)]
 
     kb = InlineKeyboardMarkup([
         [InlineKeyboardButton("⬆️ Upload File", callback_data=f"upload:{token}")],
-        [InlineKeyboardButton("❌ Cancel", callback_data="cancel_upload:0")]
+        [InlineKeyboardButton("❌ Cancel", callback_data="cancel_upload:0")],
     ])
     await message.reply_text(
         f"Found file `{filename}` in the linked message. Do you want to upload it?",
-        reply_markup=kb
+        reply_markup=kb,
     )
 
-# Override the upload callback to handle remote messages (those from message links)
-# We'll modify the upload_callback function to check if the data tuple has 4 elements instead of 3
-# and retrieve the message accordingly.
 
-@app.on_callback_query(filters.regex("^upload:"))
-async def upload_callback_v2(_, query):
-    token = query.data.replace("upload:", "")
-    data = upload_prompt_map.pop(token, None)
-    if not data:
-        return await query.answer("Upload prompt expired. Please send the file again.", show_alert=True)
+# =============================================================================
+# DIRECT FILE RECEIVER (any supported media)
+# =============================================================================
 
-    chat_id = query.message.chat.id
-    for entry in data:
-        if len(entry) == 4:
-            # Remote message: (user_id, file_chat_id, msg_id, orig_filename)
-            user_id, file_chat_id, msg_id, orig_filename = entry
-            try:
-                file_msg = await app.get_messages(file_chat_id, msg_id)
-            except Exception:
-                await query.answer(f"Message for {orig_filename} missing.", show_alert=True)
-                continue
-        else:
-            # Local message: (user_id, msg_id, orig_filename)
-            user_id, msg_id, orig_filename = entry
-            try:
-                file_msg = await app.get_messages(chat_id, msg_id)
-            except Exception:
-                await query.answer(f"Message for {orig_filename} missing.", show_alert=True)
-                continue
-
-        if not file_msg:
-            await query.answer(f"File {orig_filename} no longer available.", show_alert=True)
-            continue
-
-        job_id = int(time.time() * 1000 + hash(orig_filename))
-        active_jobs[job_id] = {"status": "queued"}
-        await upload_queue.put((user_id, chat_id, file_msg, orig_filename, job_id))
-
-    await query.message.edit_text(
-        f"📦 {len(data)} file(s) added to queue.",
-        reply_markup=None
-    )
-
-# ----------- RECEIVING ALL FILE TYPES (direct messages) -----------
 @app.on_message(
-    filters.document | filters.video | filters.audio | filters.photo |
-    filters.animation | filters.sticker | filters.video_note | filters.voice
+    filters.document | filters.video | filters.audio | filters.photo
+    | filters.animation | filters.sticker | filters.video_note | filters.voice
 )
-async def on_file(_, message):
+async def on_direct_file(_, message: Message):
     if message.document:
         filename = message.document.file_name
+        file_id = message.document.file_id
     elif message.video:
         filename = message.video.file_name or f"{message.video.file_unique_id}.mp4"
+        file_id = message.video.file_id
     elif message.audio:
         filename = message.audio.file_name or f"{message.audio.file_unique_id}.mp3"
+        file_id = message.audio.file_id
     elif message.photo:
         filename = f"{message.photo.file_unique_id}.jpg"
+        file_id = message.photo.file_id
     elif message.animation:
         filename = message.animation.file_name or f"{message.animation.file_unique_id}.mp4"
+        file_id = message.animation.file_id
     elif message.sticker:
         if message.sticker.is_animated:
             filename = f"{message.sticker.file_unique_id}.tgs"
@@ -811,27 +1061,27 @@ async def on_file(_, message):
             filename = f"{message.sticker.file_unique_id}.webm"
         else:
             filename = f"{message.sticker.file_unique_id}.webp"
+        file_id = message.sticker.file_id
     elif message.video_note:
         filename = f"{message.video_note.file_unique_id}.mp4"
+        file_id = message.video_note.file_id
     elif message.voice:
         filename = f"{message.voice.file_unique_id}.ogg"
+        file_id = message.voice.file_id
     else:
         return
 
+    # Album grouping
     if message.media_group_id and (message.photo or message.video or message.document):
         group_id = str(message.media_group_id)
-        chat_id = message.chat.id
-        user_id = message.from_user.id
-
         if group_id not in album_groups:
             album_groups[group_id] = {
-                'chat_id': chat_id,
-                'user_id': user_id,
+                'chat_id': message.chat.id,
+                'user_id': message.from_user.id,
                 'messages': [],
-                'timer': None
+                'timer': None,
             }
-
-        album_groups[group_id]['messages'].append((message.id, filename))
+        album_groups[group_id]['messages'].append((message.id, filename, file_id))
 
         if album_groups[group_id]['timer']:
             album_groups[group_id]['timer'].cancel()
@@ -843,36 +1093,41 @@ async def on_file(_, message):
         album_groups[group_id]['timer'] = asyncio.create_task(timer_callback())
         return
 
-    # Single file
+    # Single file – local entry: (user_id, msg_id, filename, file_id)
     token = make_token()
-    upload_prompt_map[token] = [(message.from_user.id, message.id, filename)]
+    upload_prompt_map[token] = [(message.from_user.id, message.id, filename, file_id)]
 
     kb = InlineKeyboardMarkup([
         [InlineKeyboardButton("⬆️ Upload File", callback_data=f"upload:{token}")],
-        [InlineKeyboardButton("❌ Cancel", callback_data="cancel_upload:0")]
+        [InlineKeyboardButton("❌ Cancel", callback_data="cancel_upload:0")],
     ])
+    await message.reply_text(f"Do you want to upload `{filename}`?", reply_markup=kb)
 
-    await message.reply_text(
-        f"Do you want to upload `{filename}`?",
-        reply_markup=kb
-    )
 
-# --------------- KEEPALIVE ---------------
+# =============================================================================
+# KEEP‑ALIVE PING
+# =============================================================================
+
 def keepalive():
     while True:
         try:
-            requests.get("http://localhost:7860/", timeout=10)
-        except Exception:
-            pass
+            requests.get(SPACE_HOST, timeout=10)
+        except Exception as e:
+            logger.warning("Keep‑alive ping failed: %s", e)
         time.sleep(KEEPALIVE_INTERVAL)
 
-# --------------- MAIN ---------------
+
+# =============================================================================
+# MAIN ENTRY POINT
+# =============================================================================
+
 async def main():
     asyncio.create_task(upload_worker())
     await app.start()
-    print("Bot started")
+    logger.info("Bot started successfully")
     await asyncio.Event().wait()
 
+
 if __name__ == "__main__":
     threading.Thread(target=run_web, daemon=True).start()
     threading.Thread(target=keepalive, daemon=True).start()