Update bot/handlers.py

bot/handlers.py

@@ -11,8 +11,12 @@ from hydrogram.types import Message, CallbackQuery
 
 from bot.config import Workers
 from bot.ui import texts
-from bot.ui.keyboards import profiles_keyboard, auth_menu_keyboard
-from bot.ui.callbacks import parse_cb, AUTH_JSON, AUTH_CI, CANCEL
+from bot.ui.keyboards import profiles_keyboard, auth_menu_keyboard, main_menu_keyboard
+from bot.ui.callbacks import (
+    parse_cb,
+    AUTH_JSON, AUTH_CI, CANCEL,
+    MENU_HELP, MENU_AUTH, MENU_PROFILES, MENU_SPEEDTEST
+)
 
 from bot.core.auth import is_owner_id, require_allowed
 from bot.core.progress import SpeedETA, human_bytes, human_eta
@@ -47,19 +51,15 @@ from bot.youtube.uploader import upload_video
 from bot.integrations.diag_extra import dns_check
 from bot.integrations.http import fetch_status
 
-# ---------- in-memory auth state (simple FSM) ----------
 _AWAIT_AUTH: Dict[int, str] = {}
 
-# uptime + speedtest cooldown
 _STARTED_AT = time.time()
 _SPEED_COOLDOWN: dict[int, float] = {}
 
 
 def _pick_login_url(j: dict) -> str:
-    """Accept login_url from either top-level OR nested {data:{...}} shapes."""
     if not isinstance(j, dict):
         return ""
-
     data = j.get("data") if isinstance(j.get("data"), dict) else {}
 
     def _get(d: dict, k: str) -> str:
@@ -78,14 +78,54 @@ def _pick_login_url(j: dict) -> str:
     )
 
 
-def _fmt_gb(b: float) -> float:
-    return float(b) / (1024.0**3)
+async def _render_speedtest(mb: int) -> str:
+    mb = max(1, min(200, int(mb)))
+    dl_bytes = mb * 1024 * 1024
+    up_mb = min(20, max(3, mb // 5))  # 3..20MB
+    up_bytes = up_mb * 1024 * 1024
+
+    up_sec = int(max(0, time.time() - _STARTED_AT))
+    up_txt = human_eta(up_sec)
+
+    d = disk_total_free("/")
+    total_gb = d["total"] / (1024.0**3)
+    free_gb = d["free"] / (1024.0**3)
+
+    dl_task = asyncio.create_task(net_download_test(bytes_target=dl_bytes, timeout=40.0))
+    up_task = asyncio.create_task(net_upload_test(bytes_target=up_bytes, timeout=40.0))
+    p_task = asyncio.create_task(ping_ms())
+
+    dl, up, p = await asyncio.gather(dl_task, up_task, p_task)
+
+    ping_line = f"{p:.0f} ms" if p is not None else "N/A"
+    dl_mb_s = bytes_per_sec_to_mb_s(dl["bps"])
+    up_mb_s = bytes_per_sec_to_mb_s(up["bps"])
+
+    return (
+        "⚡ **Server Speed Test**\n\n"
+        "🕒 Uptime\n"
+        f"└ {up_txt}\n\n"
+        "📡 Network\n"
+        f"├ 🟢 Ping: {ping_line}\n"
+        f"├ ⬇️ Download: {dl_mb_s:.2f} MB/s\n"
+        f"│   └ {bytes_to_mb(dl['bytes']):.2f} MB in {dl['seconds']:.2f}s\n"
+        f"└ ⬆️ Upload: {up_mb_s:.2f} MB/s\n"
+        f"    └ {bytes_to_mb(up['bytes']):.2f} MB in {up['seconds']:.2f}s\n\n"
+        "💾 Storage\n"
+        f"├ Total: {total_gb:.1f} GB\n"
+        f"└ Free : {free_gb:.1f} GB\n\n"
+        "🧪 Measured from server (real MB/s, not Mbps)"
+    )
 
 
 def setup_handlers(app: Client) -> None:
     @app.on_message(filters.command(["start"]) & filters.private)
     async def start_handler(_: Client, m: Message):
-        await safe_reply(m, texts.START_TEXT)
+        await safe_reply(m, texts.START_TEXT, reply_markup=main_menu_keyboard())
+
+    @app.on_message(filters.command(["help"]) & filters.private)
+    async def help_handler(_: Client, m: Message):
+        await safe_reply(m, texts.HELP_TEXT, reply_markup=main_menu_keyboard())
 
     @app.on_message(filters.command(["ping"]) & filters.private)
     async def ping_handler(_: Client, m: Message):
@@ -107,50 +147,20 @@ def setup_handlers(app: Client) -> None:
             return await safe_reply(m, "⏳ Wait 30s then try again.")
         _SPEED_COOLDOWN[uid] = now
 
-        msg = await safe_reply(m, "⚡ Running speed test…")
-
-        # uptime
-        up_sec = int(max(0, time.time() - _STARTED_AT))
-        up_txt = human_eta(up_sec)
-
-        # disk
-        d = disk_total_free("/")
-        total_gb = _fmt_gb(d["total"])
-        free_gb = _fmt_gb(d["free"])
-
-        # run parallel
-        dl_task = asyncio.create_task(net_download_test())
-        up_task = asyncio.create_task(net_upload_test())
-        p_task = asyncio.create_task(ping_ms())
-
-        dl, up, p = await asyncio.gather(dl_task, up_task, p_task)
-
-        ping_line = f"{p:.0f} ms" if p is not None else "N/A"
-
-        dl_mb_s = bytes_per_sec_to_mb_s(dl["bps"])
-        up_mb_s = bytes_per_sec_to_mb_s(up["bps"])
-
-        dl_mb = bytes_to_mb(dl["bytes"])
-        up_mb = bytes_to_mb(up["bytes"])
-
-        txt = (
-            "⚡ Server Speed Test\n\n"
-            "🕒 Uptime\n"
-            f"└ {up_txt}\n\n"
-            "📡 Network\n"
-            f"├ 🟢 Ping: {ping_line}\n"
-            f"├ ⬇️ Download: {dl_mb_s:.2f} MB/s\n"
-            f"│   └ {dl_mb:.2f} MB in {dl['seconds']:.2f}s\n"
-            f"└ ⬆️ Upload: {up_mb_s:.2f} MB/s\n"
-            f"    └ {up_mb:.2f} MB in {up['seconds']:.2f}s\n\n"
-            "💾 Storage\n"
-            f"├ Total: {total_gb:.1f} GB\n"
-            f"└ Free : {free_gb:.1f} GB\n\n"
-            "🧪 Measured from server (real MB/s, not Mbps)"
-        )
-        await safe_edit(msg, txt)
+        parts = (m.text or "").split()
+        try:
+            mb = int(parts[1]) if len(parts) > 1 else 8
+        except Exception:
+            mb = 8
+
+        msg = await safe_reply(m, "⚡ Running speedtest…")
+        try:
+            txt = await _render_speedtest(mb)
+            await safe_edit(msg, txt)
+        except Exception as e:
+            await safe_edit(msg, f"❌ Speedtest failed: `{type(e).__name__}: {e}`")
 
-    # -------- OWNER allow/disallow via forward OR id --------
+    # -------- OWNER allow/disallow --------
     @app.on_message(filters.command(["allow"]) & filters.private)
     async def allow_cmd(_: Client, m: Message):
         if not is_owner_id(m.from_user.id if m.from_user else None):
@@ -184,7 +194,6 @@ def setup_handlers(app: Client) -> None:
         j = await disallow_user(target)
         await safe_reply(m, f"✅ disallowed: `{target}`\n`{j}`")
 
-    # -------- Owner stats + diag --------
     @app.on_message(filters.command(["stats"]) & filters.private)
     async def stats_cmd(_: Client, m: Message):
         if not is_owner_id(m.from_user.id if m.from_user else None):
@@ -210,38 +219,64 @@ def setup_handlers(app: Client) -> None:
             f"DNS google -> {dns_check('https://www.google.com')}\n",
         )
 
-    # -------- AUTH UI --------
+    # -------- AUTH --------
     @app.on_message(filters.command(["auth"]) & filters.private)
     async def auth_cmd(_: Client, m: Message):
         uid = m.from_user.id if m.from_user else 0
-        ok = await require_allowed(uid)
-        if not ok:
+        if not await require_allowed(uid):
             return await safe_reply(m, texts.NOT_ALLOWED)
 
         _AWAIT_AUTH.pop(uid, None)
-        await safe_reply(m, texts.AUTH_MENU, reply_markup=auth_menu_keyboard())
+        await safe_reply(m, "🔐 Add YouTube Profile\n\nChoose method:", reply_markup=auth_menu_keyboard())
 
     @app.on_message(filters.command(["cancel"]) & filters.private)
     async def cancel_cmd(_: Client, m: Message):
         _AWAIT_AUTH.pop(m.from_user.id, None)
-        await safe_reply(m, texts.CANCELLED)
+        await safe_reply(m, texts.CANCELLED, reply_markup=main_menu_keyboard())
 
     @app.on_callback_query()
     async def cb(_: Client, q: CallbackQuery):
         uid = q.from_user.id if q.from_user else 0
         action, value = parse_cb(q.data or "")
 
+        # main menu
+        if action == MENU_HELP:
+            try:
+                await q.message.edit_text(texts.HELP_TEXT, reply_markup=main_menu_keyboard())
+            except Exception:
+                pass
+            return await q.answer("Help")
+
+        if action == MENU_SPEEDTEST:
+            await q.answer("Use /speedtest", show_alert=False)
+            return
+
+        if action == MENU_AUTH:
+            if not await require_allowed(uid):
+                return await q.answer("Not allowed", show_alert=True)
+            _AWAIT_AUTH.pop(uid, None)
+            try:
+                await q.message.edit_text("🔐 Add YouTube Profile\n\nChoose method:", reply_markup=auth_menu_keyboard())
+            except Exception:
+                pass
+            return await q.answer("Auth")
+
+        if action == MENU_PROFILES:
+            await q.answer("Use /profiles", show_alert=False)
+            return
+
+        # auth menu
         if action in (AUTH_JSON, AUTH_CI, CANCEL):
             if action == CANCEL:
                 _AWAIT_AUTH.pop(uid, None)
                 return await q.answer("Cancelled")
             if action == AUTH_JSON:
                 _AWAIT_AUTH[uid] = "json"
-                await q.message.edit_text(texts.ASK_JSON)
+                await q.message.edit_text("📄 Paste your client JSON here.\n\n(You can /cancel anytime)")
                 return await q.answer("Paste JSON")
             if action == AUTH_CI:
                 _AWAIT_AUTH[uid] = "ci"
-                await q.message.edit_text(texts.ASK_ID_SECRET)
+                await q.message.edit_text("🔑 Send Client ID and Client Secret.\n\nFormat:\n<id> | <secret>\n(or 2 lines)\n\n(/cancel anytime)")
                 return await q.answer("Send ID & Secret")
             return
 
@@ -255,32 +290,18 @@ def setup_handlers(app: Client) -> None:
 
             await q.answer("Current set ✅", show_alert=False)
             try:
-                await q.message.edit_text(f"✅ Current profile set: `{value}`")
+                await q.message.edit_text(f"🟢 Current profile set:\n`{value}`", reply_markup=main_menu_keyboard())
             except Exception:
                 pass
             return
 
         await q.answer("OK")
 
-    # message sink for auth states
+    # auth sink
     @app.on_message(
         filters.private
         & filters.text
-        & ~filters.command(
-            [
-                "start",
-                "ping",
-                "me",
-                "auth",
-                "cancel",
-                "profiles",
-                "speedtest",
-                "stats",
-                "allow",
-                "disallow",
-                "diag",
-            ]
-        )
+        & ~filters.command(["start", "help", "ping", "me", "auth", "cancel", "profiles", "speedtest", "stats", "allow", "disallow", "diag"])
     )
     async def auth_sink(_: Client, m: Message):
         uid = m.from_user.id if m.from_user else 0
@@ -304,7 +325,6 @@ def setup_handlers(app: Client) -> None:
                 client_id, client_secret = parts[0], parts[1]
 
             j = await profile_add(uid, client_id, client_secret, label="main", ttl_sec=600)
-
             if not isinstance(j, dict):
                 return await safe_reply(m, f"❌ profile_add bad response: `{j}`")
 
@@ -316,7 +336,7 @@ def setup_handlers(app: Client) -> None:
                 return await safe_reply(m, f"❌ profile_add OK but login link missing.\nBackend response:\n`{j}`")
 
             _AWAIT_AUTH.pop(uid, None)
-            await safe_reply(m, texts.SENT_AUTH_LINK + login_url)
+            await safe_reply(m, texts.SENT_AUTH_LINK + login_url, reply_markup=main_menu_keyboard())
 
         except Exception as e:
             await safe_reply(m, f"{texts.PARSE_FAIL}\n`{type(e).__name__}: {e}`")
@@ -324,8 +344,7 @@ def setup_handlers(app: Client) -> None:
     @app.on_message(filters.command(["profiles"]) & filters.private)
     async def profiles_cmd(_: Client, m: Message):
         uid = m.from_user.id if m.from_user else 0
-        ok = await require_allowed(uid)
-        if not ok:
+        if not await require_allowed(uid):
             return await safe_reply(m, texts.NOT_ALLOWED)
 
         j = await profile_list(uid)
@@ -335,7 +354,9 @@ def setup_handlers(app: Client) -> None:
         profiles = j.get("profiles") or []
         default_id = j.get("default_profile_id")
 
-        txt = f"Current: `{default_id}`\n\n"
+        txt = "👤 **Profiles**\n\n"
+        txt += f"🟢 Current: `{default_id}`\n\n"
+
         for i, p in enumerate(profiles, start=1):
             pid = p.get("profile_id")
             ch = (p.get("channel_title") or "—").strip()
@@ -343,17 +364,19 @@ def setup_handlers(app: Client) -> None:
             sec = p.get("client_secret_hint") or "****…****"
             connected = "✅ Connected" if (p.get("has_refresh") and p.get("channel_id")) else "⏳ Pending"
             cur = "🟢 Current" if pid == default_id else ""
-            txt += f"[{i}] {ch} | {cur} | id: {cid} | secret: {sec} | {connected}\n"
+            txt += f"[{i}] {ch}  {cur}\n"
+            txt += f"    id: {cid}\n"
+            txt += f"    secret: {sec}\n"
+            txt += f"    {connected}\n\n"
 
         kb = profiles_keyboard(profiles)
         await safe_reply(m, txt, reply_markup=kb)
 
-    # -------- Upload handler --------
+    # -------- Upload --------
     @app.on_message(filters.private & (filters.video | filters.document))
     async def upload_handler(app_: Client, m: Message):
         uid = m.from_user.id if m.from_user else 0
-        ok = await require_allowed(uid)
-        if not ok:
+        if not await require_allowed(uid):
             return await safe_reply(m, texts.NOT_ALLOWED)
 
         media = m.video or m.document
@@ -384,8 +407,7 @@ def setup_handlers(app: Client) -> None:
             return await safe_reply(m, texts.TOKEN_FAIL.format(tok))
         access_tok = tok["access_token"]
 
-        status = await safe_reply(m, texts.UPLOAD_START)
-
+        status = await safe_reply(m, "⬇️ Downloading…")
         task_id = f"{uid}:{m.id}"
         create_task(task_id, uid)
 
@@ -394,15 +416,13 @@ def setup_handlers(app: Client) -> None:
         file_name = "video.bin"
 
         try:
-            # ---------------- Download (Telegram) ----------------
+            # ---- Download ----
             set_task(task_id, "downloading", "")
-
             dl_t0 = time.time()
             file_path, file_size, file_name = await download_to_temp(app_, m)
             dl_sec = max(0.001, time.time() - dl_t0)
 
-            # fallback size from disk if TG didn't give it
-            if not file_size and file_path and os.path.exists(file_path):
+            if (not file_size) and file_path and os.path.exists(file_path):
                 try:
                     file_size = os.path.getsize(file_path)
                 except Exception:
@@ -410,10 +430,17 @@ def setup_handlers(app: Client) -> None:
 
             dl_avg_bps = int((file_size or 0) / dl_sec) if file_size else 0
 
-            # parse title/desc after download
-            title, desc = extract_title_description(m, file_name)
+            await safe_edit(
+                status,
+                "✅ Downloaded\n"
+                f"• File: `{file_name}`\n"
+                f"• Size: {human_bytes(file_size)}\n"
+                f"• Time: {human_eta(int(dl_sec))}\n"
+                f"• Avg: {human_bytes(dl_avg_bps)}/s\n\n"
+                "⬆️ Starting upload…",
+            )
 
-            # ---------------- Upload (YouTube) ----------------
+            # ---- Upload ----
             set_task(task_id, "uploading", "")
             up_t0 = time.time()
 
@@ -427,12 +454,14 @@ def setup_handlers(app: Client) -> None:
                     return
                 last_edit = time.time()
                 txt = (
-                    f"{texts.UPLOAD_TO_YT}\n"
+                    "⬆️ Uploading to YouTube…\n"
                     f"{human_bytes(done)}/{human_bytes(total)}\n"
                     f"speed: {human_bytes(snap['speed_bps'])}/s | eta: {human_eta(snap['eta_sec'])}"
                 )
                 await safe_edit(status, txt)
 
+            title, desc = extract_title_description(m, file_name)
+
             yt_url = await upload_video(
                 access_tok,
                 file_path,
@@ -446,11 +475,11 @@ def setup_handlers(app: Client) -> None:
             up_avg_bps = int((file_size or 0) / up_sec) if file_size else 0
 
             await record_upload(uid, use_profile_id)
-
             set_task(task_id, "done", "", yt_url=yt_url)
-            await safe_reply(
-                m,
-                f"{texts.DONE}\n\n"
+
+            await safe_edit(
+                status,
+                "✅ **Uploaded!**\n\n"
                 f"🎬 {title}\n"
                 f"🔗 {yt_url}\n\n"
                 f"⬇️ Download: {human_eta(int(dl_sec))} | avg {human_bytes(dl_avg_bps)}/s\n"
@@ -459,7 +488,7 @@ def setup_handlers(app: Client) -> None:
 
         except Exception as e:
             set_task(task_id, "error", str(e))
-            await safe_reply(m, f"❌ Upload failed: `{type(e).__name__}: {e}`")
+            await safe_edit(status, f"❌ Upload failed:\n`{type(e).__name__}: {e}`")
         finally:
             if file_path:
                 cleanup_file(file_path)