Update main.py

main.py

@@ -1,15 +1,9 @@
 # Link Generator Bot
 # Webhook-only | HuggingFace Spaces compatible
-#
-# Flow:
-#   User sends any file/photo/video/audio to the bot
-#   → file_id extracted from raw webhook update
-#   → Downloaded via Pyrogram MTProto (bot.download_media)
-#   → Saved to HF Space local storage
-#   → Public link returned to user
 
 import os
 import sys
+import time
 import logging
 import asyncio
 import traceback
@@ -22,7 +16,6 @@ from fastapi.responses import FileResponse, JSONResponse
 from pyrogram import Client
 from pyrogram.enums import ParseMode
 from pyrogram.errors import FloodWait
-from pyrogram.types import Message
 
 from fileManager import FileManager
 
@@ -111,6 +104,51 @@ def track_task(coro):
     task.add_done_callback(RUNNING_TASKS.discard)
     return task
 
+# ============================================================
+# UX helpers
+# ============================================================
+def fmt_size(b: float) -> str:
+    for u in ["B", "KB", "MB", "GB"]:
+        if b < 1024:
+            return f"{b:.1f} {u}"
+        b /= 1024
+    return f"{b:.1f} TB"
+
+def fmt_time(seconds: float) -> str:
+    seconds = max(0, int(seconds))
+    if seconds < 60:
+        return f"{seconds}s"
+    m, s = divmod(seconds, 60)
+    if m < 60:
+        return f"{m}m {s}s"
+    h, m = divmod(m, 60)
+    return f"{h}h {m}m"
+
+def progress_bar(current: int, total: int, length: int = 10) -> str:
+    if total <= 0:
+        return "▓" * length
+    filled = int(length * current / total)
+    empty  = length - filled
+    return "▓" * filled + "░" * empty
+
+def build_progress_text(phase: str, current: int, total: int,
+                        speed: float, elapsed: float, filename: str) -> str:
+    pct   = min(100, int(current * 100 / total)) if total > 0 else 0
+    bar   = progress_bar(current, total)
+    eta   = fmt_time((total - current) / speed) if speed > 0 else "—"
+    icons = {"download": "📥", "upload": "📤"}
+    icon  = icons.get(phase, "⚙️")
+    label = "Downloading" if phase == "download" else "Saving"
+
+    return (
+        f"{icon} **{label}:** `{filename}`\n\n"
+        f"`{bar}` **{pct}%**\n\n"
+        f"**Done:** `{fmt_size(current)}` / `{fmt_size(total)}`\n"
+        f"**Speed:** `{fmt_size(speed)}/s`\n"
+        f"**ETA:** `{eta}`\n"
+        f"**Elapsed:** `{fmt_time(elapsed)}`"
+    )
+
 # ============================================================
 # Webhook response helpers
 # ============================================================
@@ -153,110 +191,167 @@ def handle_help(chat_id: int) -> dict:
 
 def handle_stats(chat_id: int) -> dict:
     import shutil, psutil
+    from time import time as _time
     total, used, free = shutil.disk_usage(".")
     proc = psutil.Process(os.getpid())
-    def fmt(b):
-        for u in ["B", "KB", "MB", "GB"]:
-            if b < 1024: return f"{b:.1f} {u}"
-            b /= 1024
-        return f"{b:.1f} TB"
-    fl_count = len(os.listdir("fl")) if os.path.isdir("fl") else 0
+    fl_count = len([f for f in os.listdir("fl")
+                    if not f.startswith("tmp_")]) if os.path.isdir("fl") else 0
+    active   = len([t for t in RUNNING_TASKS if not t.done()])
     return _msg(chat_id,
-        "**📊 Bot Stats**\n\n"
-        f"**➜ Disk:** `{fmt(used)}` used / `{fmt(total)}` total\n"
-        f"**➜ Free:** `{fmt(free)}`\n"
-        f"**➜ Memory:** `{round(proc.memory_info()[0]/1024**2)} MiB`\n"
-        f"**➜ CPU:** `{psutil.cpu_percent(interval=0.2)}%`\n"
-        f"**➜ Stored files:** `{fl_count}`")
+        "📊 **Bot Statistics**\n"
+        "━━━━━━━━━━━━━━━━━━━\n"
+        f"💾 **Disk Used:**    `{fmt_size(used)}` / `{fmt_size(total)}`\n"
+        f"🆓 **Disk Free:**    `{fmt_size(free)}`\n"
+        f"🧠 **Memory:**       `{round(proc.memory_info()[0]/1024**2)} MiB`\n"
+        f"⚡ **CPU:**          `{psutil.cpu_percent(interval=0.2)}%`\n"
+        f"📁 **Stored Files:** `{fl_count}`\n"
+        f"⚙️ **Active Jobs:**  `{active}`")
 
 # ============================================================
-# Extract file_id + filename from raw webhook message dict
+# Extract file_id + filename + size from raw webhook message
 # ============================================================
 MEDIA_TYPES = [
     "document", "video", "audio", "voice",
     "photo", "animation", "sticker", "video_note"
 ]
-
 EXT_MAP = {
     "video": "mp4", "audio": "mp3", "voice": "ogg",
     "animation": "gif", "sticker": "webp", "video_note": "mp4"
 }
 
 def get_file_info(msg: dict) -> tuple:
-    """Returns (file_id, filename) or (None, None)."""
+    """Returns (file_id, filename, file_size) or (None, None, 0)."""
     for mtype in MEDIA_TYPES:
         if mtype not in msg:
             continue
         if mtype == "photo":
-            # list of sizes — pick the largest
             photos = msg["photo"]
             obj    = photos[-1] if isinstance(photos, list) else photos
-            return obj["file_id"], f"{obj.get('file_unique_id', 'photo')}.jpg"
-
+            return obj["file_id"], f"{obj.get('file_unique_id', 'photo')}.jpg", obj.get("file_size", 0)
         obj   = msg[mtype]
         fid   = obj["file_id"]
         fname = obj.get("file_name") or obj.get("file_unique_id", mtype)
         if "." not in fname.split("/")[-1]:
             fname = f"{fname}.{EXT_MAP.get(mtype, 'bin')}"
-        return fid, fname
-
-    return None, None
+        return fid, fname, obj.get("file_size", 0)
+    return None, None, 0
 
 # ============================================================
-# Async media handler — uses file_id directly
+# Async media handler with live progress
 # ============================================================
-async def handle_media(chat_id: int, reply_to_id: int, file_id: str, filename: str):
-    progress_msg = await bot.send_message(
-        chat_id, "⏳ **Processing your file...**",
-        reply_to_message_id=reply_to_id)
-    try:
-        tmp_path = f"fl/tmp_{reply_to_id}_{filename}"
+async def handle_media(chat_id: int, reply_to_id: int,
+                       file_id: str, filename: str, file_size: int):
+
+    short_name = filename if len(filename) <= 30 else filename[:27] + "..."
+
+    # ── Phase 0: initial status ──────────────────────────────
+    status_msg = await bot.send_message(
+        chat_id,
+        f"🔍 **Preparing to process your file…**\n\n"
+        f"📄 `{short_name}`\n"
+        f"📦 Size: `{fmt_size(file_size) if file_size else 'unknown'}`",
+        reply_to_message_id=reply_to_id,
+    )
+
+    dl_start   = time.time()
+    last_edit  = [0.0]   # mutable for closure
+    tmp_path   = f"fl/tmp_{reply_to_id}_{filename}"
 
-        # Download via Pyrogram MTProto using file_id directly
+    # ── Progress callback for Pyrogram download ──────────────
+    async def on_progress(current: int, total: int):
+        now = time.time()
+        if now - last_edit[0] < 2.5:          # throttle edits to every 2.5s
+            return
+        last_edit[0] = now
+        elapsed = now - dl_start
+        speed   = current / elapsed if elapsed > 0 else 0
+        try:
+            await status_msg.edit(
+                build_progress_text("download", current, total,
+                                    speed, elapsed, short_name)
+            )
+        except Exception:
+            pass
+
+    try:
+        # ── Phase 1: Download ────────────────────────────────
+        downloaded = None
         for attempt in range(2):
             try:
-                downloaded = await bot.download_media(file_id, file_name=tmp_path)
+                downloaded = await bot.download_media(
+                    file_id,
+                    file_name=tmp_path,
+                    progress=on_progress,
+                )
                 break
             except FloodWait as e:
                 wait_s = int(getattr(e, "value", 0) or 0)
                 LOGGER(__name__).warning(f"FloodWait: {wait_s}s")
                 if attempt == 0 and wait_s > 0:
+                    await status_msg.edit(
+                        f"⏳ **FloodWait — retrying in {wait_s}s…**\n\n"
+                        f"📄 `{short_name}`"
+                    )
                     await asyncio.sleep(wait_s + 1)
+                    dl_start = time.time()
                     continue
                 raise
 
-        if not downloaded or not os.path.exists(downloaded) or os.path.getsize(downloaded) == 0:
-            await progress_msg.edit("❌ **Download failed. Please try again.**")
+        if not downloaded or not os.path.exists(downloaded) \
+                or os.path.getsize(downloaded) == 0:
+            await status_msg.edit(
+                "❌ **Download Failed**\n\n"
+                "Could not retrieve the file from Telegram.\n"
+                "Please try again."
+            )
             try: os.remove(tmp_path)
             except: pass
             return
 
-        # Read bytes → save via FileManager → get public URL
+        actual_size = os.path.getsize(downloaded)
+        dl_elapsed  = time.time() - dl_start
+        dl_speed    = actual_size / dl_elapsed if dl_elapsed > 0 else 0
+
+        # ── Phase 2: Saving ──────────────────────────────────
+        await status_msg.edit(
+            f"📤 **Saving to server…**\n\n"
+            f"📄 `{short_name}`\n"
+            f"📦 `{fmt_size(actual_size)}`\n\n"
+            f"⬇️ Downloaded in `{fmt_time(dl_elapsed)}` "
+            f"at `{fmt_size(dl_speed)}/s`"
+        )
+
         with open(downloaded, "rb") as f:
             file_data = f.read()
         os.remove(downloaded)
 
-        entry     = fm.save_file(file_data, filename)
-        # Build direct download link served by our own FastAPI endpoint
-        link = f"{SPACE_URL}/download/{entry['id']}"
-        file_size = len(file_data)
-        size_str  = (f"{file_size/1024/1024:.2f} MB"
-                     if file_size > 1024 * 1024
-                     else f"{file_size/1024:.1f} KB")
+        entry = fm.save_file(file_data, filename)
+        link  = f"{SPACE_URL}/download/{entry['id']}"
 
-        await progress_msg.edit(
-            f"✅ **Link Generated!**\n\n"
+        # ── Phase 3: Done ────────────────────────────────────
+        total_elapsed = time.time() - dl_start
+        await status_msg.edit(
+            f"✅ **Link Generated Successfully!**\n"
+            f"━━━━━━━━━━━━━━━━━━━\n"
             f"🔗 **Link:**\n`{link}`\n\n"
             f"📄 **File:** `{entry['filename']}`\n"
-            f"📦 **Size:** `{size_str}`\n\n"
-            f"_Link is public and accessible by anyone._"
+            f"📦 **Size:** `{fmt_size(actual_size)}`\n"
+            f"⚡ **Speed:** `{fmt_size(dl_speed)}/s`\n"
+            f"⏱ **Time:** `{fmt_time(total_elapsed)}`\n"
+            f"━━━━━━━━━━━━━━━━━━━\n"
+            f"_Tap the link to copy • Accessible by anyone_",
+            disable_web_page_preview=True,
         )
-        LOGGER(__name__).info(f"Generated link for {filename}: {link}")
+        LOGGER(__name__).info(f"Generated link for {filename} ({fmt_size(actual_size)}): {link}")
 
     except Exception as e:
         LOGGER(__name__).error(traceback.format_exc())
         try:
-            await progress_msg.edit(f"❌ **Error:** `{str(e)}`")
+            await status_msg.edit(
+                f"❌ **Something went wrong**\n\n"
+                f"`{str(e)[:200]}`\n\n"
+                f"Please try again or contact support."
+            )
         except Exception:
             pass
 
@@ -267,7 +362,8 @@ async def handle_media(chat_id: int, reply_to_id: int, file_id: str, filename: s
 async def lifespan(app: FastAPI):
     LOGGER(__name__).info("Starting bot client…")
     await bot.start()
-    LOGGER(__name__).info("Bot ready — waiting for webhook updates.")
+    me = await bot.get_me()
+    LOGGER(__name__).info(f"Bot ready: @{me.username} — waiting for webhook updates.")
     yield
     LOGGER(__name__).info("Shutting down…")
     await bot.stop()
@@ -281,10 +377,8 @@ app = FastAPI(lifespan=lifespan)
 async def root():
     return {"status": "ok", "service": "Link Generator Bot"}
 
-
 @app.get("/download/{file_id}")
 async def download_file(file_id: str):
-    """Serve a stored file by its UUID."""
     entry = fm.get_entry_by_id(file_id)
     if not entry:
         return JSONResponse({"error": "File not found or expired"}, status_code=404)
@@ -313,7 +407,6 @@ async def telegram_webhook(request: Request):
     text    = msg.get("text", "").strip()
     msg_id  = msg["message_id"]
 
-    # ── commands ──────────────────────────────────────────
     if text.startswith("/start"):
         return JSONResponse(handle_start(chat_id))
     if text.startswith("/help"):
@@ -321,13 +414,11 @@ async def telegram_webhook(request: Request):
     if text.startswith("/stats"):
         return JSONResponse(handle_stats(chat_id))
 
-    # ── media ─────────────────────────────────────────────
-    file_id, filename = get_file_info(msg)
+    file_id, filename, file_size = get_file_info(msg)
     if file_id:
-        track_task(handle_media(chat_id, msg_id, file_id, filename))
+        track_task(handle_media(chat_id, msg_id, file_id, filename, file_size))
         return JSONResponse({"status": "ok"})
 
-    # ── unknown text ──────────────────────────────────────
     if text and not text.startswith("/"):
         return JSONResponse(_msg(chat_id,
             "📎 **Please send a file, photo, video or audio.**\n"