Update app.py

app.py

@@ -80,7 +80,7 @@ def db_remove_file(enc_path):
     cursor.execute("DELETE FROM files WHERE enc_path = ?", (enc_path,))
     conn.commit()
 
-# --------------- CALLBACK MAPPINGS (avoid long data) ---------------
+# --------------- CALLBACK MAPPINGS ---------------
 file_callback_map = {}     # token -> enc_path
 upload_prompt_map = {}     # token -> (user_id, message_id, orig_filename)
 
@@ -136,14 +136,8 @@ def format_size(bytes_val):
 def format_speed(bytes_per_sec):
     return f"{format_size(bytes_per_sec)}/s"
 
-# --------------- PROGRESS UPDATER (GENERIC) ---------------
-def start_progress_updater(job_id: int, status_msg, cancel_event: asyncio.Event,
-                           total_bytes: Optional[int] = None, action: str = "uploading"):
-    """
-    Launches a background task that edits status_msg every 3 seconds with progress.
-    action is a string like "uploading", "downloading", "sending"
-    Returns {'task': asyncio.Task, 'bytes': list} where bytes[0] should be updated regularly.
-    """
+async def send_progress_updater(job_id, status_msg, cancel_event, total_bytes):
+    """Shows Telegram send progress every 3s."""
     bytes_holder = [0]
     start_time = time.time()
 
@@ -152,15 +146,8 @@ def start_progress_updater(job_id: int, status_msg, cancel_event: asyncio.Event,
             elapsed = max(time.time() - start_time, 0.01)
             current = bytes_holder[0]
             speed = current / elapsed
-            text = ""
-            if total_bytes and total_bytes > 0:
-                pct = min(100 * current / total_bytes, 100)
-                text = f"{pct:.1f}% of {format_size(total_bytes)}\n"
-            text += f"{format_size(current)} {action}\n"
-            text += f"Speed: {format_speed(speed)}\n"
-            if total_bytes and speed > 0:
-                eta = (total_bytes - current) / speed
-                text += f"ETA: {int(eta)}s"
+            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"
             try:
                 await status_msg.edit_text(
                     text,
@@ -171,21 +158,16 @@ def start_progress_updater(job_id: int, status_msg, cancel_event: asyncio.Event,
             except Exception:
                 break
             await asyncio.sleep(3)
-        try:
-            await status_msg.edit_text("❌ Cancelled.", reply_markup=None)
-        except Exception:
-            pass
 
     task = asyncio.create_task(updater())
-    return {'task': task, 'bytes': bytes_holder}
+    return task, bytes_holder
 
 # --------------- UPLOAD JOB ---------------
-async def do_upload_job(user_id: int, chat_id: int, file_msg, orig_filename: str,
-                        job_id: int, cancel_event: asyncio.Event, status_msg):
+async def do_upload_job(user_id, chat_id, file_msg, orig_filename, job_id, cancel_event, status_msg):
     temp_dir = tempfile.mkdtemp()
     local_path = None
     try:
-        # Step 1: Download from Telegram with progress
+        # 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}")]
@@ -218,29 +200,18 @@ async def do_upload_job(user_id: int, chat_id: int, file_msg, orig_filename: str
         )
         if cancel_event.is_set():
             return
+
         file_size = os.path.getsize(local_path)
 
-        # Step 2: Upload to Hugging Face (no real-time progress from library)
+        # 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}"
 
-        prog = start_progress_updater(job_id, status_msg, cancel_event,
-                                      total_bytes=file_size, action="uploaded")
-        upload_prog_task = prog['task']
-        upload_bytes_holder = prog['bytes']
-
-        # Simulate progress based on time (assume ~10 MB/s)
-        async def simulate_upload():
-            start = time.time()
-            while not cancel_event.is_set():
-                elapsed = time.time() - start
-                estimated = min(elapsed * 10_000_000, file_size)
-                upload_bytes_holder[0] = estimated
-                await asyncio.sleep(1)
-            upload_bytes_holder[0] = file_size
-
-        sim_task = asyncio.create_task(simulate_upload())
+        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()
@@ -261,7 +232,6 @@ async def do_upload_job(user_id: int, chat_id: int, file_msg, orig_filename: str
             await asyncio.sleep(1)
 
         if cancel_event.is_set():
-            # Try to delete the partially uploaded file
             try:
                 hf.delete_file(
                     path_in_repo=repo_path,
@@ -270,17 +240,10 @@ async def do_upload_job(user_id: int, chat_id: int, file_msg, orig_filename: str
                 )
             except Exception:
                 pass
-            upload_prog_task.cancel()
-            sim_task.cancel()
             await status_msg.edit_text("❌ Upload cancelled.", reply_markup=None)
             return
 
         await task
-        upload_bytes_holder[0] = file_size
-        await asyncio.sleep(0.5)
-        upload_prog_task.cancel()
-        sim_task.cancel()
-
         db_add_file(user_id, repo_path, orig_filename, file_size)
         await status_msg.edit_text(f"✅ Uploaded: `{orig_filename}`", reply_markup=None)
 
@@ -290,13 +253,12 @@ async def do_upload_job(user_id: int, chat_id: int, file_msg, orig_filename: str
         cleanup(temp_dir)
 
 # --------------- DOWNLOAD JOB ---------------
-async def do_download_job(chat_id: int, enc_path: str, orig_name: str,
-                          job_id: int, cancel_event: asyncio.Event, status_msg):
+async def do_download_job(chat_id, enc_path, orig_name, job_id, cancel_event, status_msg):
     temp_dir = tempfile.mkdtemp()
     try:
         loop = asyncio.get_event_loop()
 
-        # Get file metadata
+        # 1. Get file size
         file_url = hf_hub_url(repo_id=DATASET_REPO, filename=enc_path, repo_type="dataset")
         metadata = await loop.run_in_executor(
             None,
@@ -304,29 +266,28 @@ async def do_download_job(chat_id: int, enc_path: str, orig_name: str,
         )
         total_size = metadata.size
 
-        # Progress for download
-        prog = start_progress_updater(job_id, status_msg, cancel_event,
-                                      total_bytes=total_size, action="downloaded")
-        prog_task = prog['task']
-        bytes_holder = prog['bytes']
+        # 2. Download from HF (no speed/progress)
+        await status_msg.edit_text("⬇️ Downloading from dataset...",
+                                   reply_markup=InlineKeyboardMarkup([
+                                       [InlineKeyboardButton("❌ Cancel", callback_data=f"cancel_job:{job_id}")]
+                                   ]))
 
         local_file = os.path.join(temp_dir, "downloaded_file")
 
-        def sync_download_with_progress():
+        download_done = loop.create_future()
+
+        def sync_download():
             headers = {"Authorization": f"Bearer {HF_TOKEN}"}
             resp = requests.get(file_url, stream=True, headers=headers)
             resp.raise_for_status()
             with open(local_file, "wb") as f:
-                downloaded = 0
                 for chunk in resp.iter_content(chunk_size=8192):
                     if cancel_event.is_set():
                         break
                     f.write(chunk)
-                    downloaded += len(chunk)
-                    bytes_holder[0] = downloaded
             return True
 
-        task = loop.run_in_executor(None, sync_download_with_progress)
+        task = loop.run_in_executor(None, sync_download)
 
         while not task.done():
             if cancel_event.is_set():
@@ -334,27 +295,20 @@ async def do_download_job(chat_id: int, enc_path: str, orig_name: str,
             await asyncio.sleep(1)
 
         if cancel_event.is_set():
-            prog_task.cancel()
             await status_msg.edit_text("❌ Download cancelled.", reply_markup=None)
             return
 
         await task
-        bytes_holder[0] = total_size
-        await asyncio.sleep(0.5)
-        prog_task.cancel()
 
+        # 3. Send back to Telegram (with speed/progress if single file)
         file_size = os.path.getsize(local_file)
-
-        # Send file (split if needed) with upload progress
         if file_size <= MAX_FILE_SIZE_TG:
-            # Progress for sending
-            prog_send = start_progress_updater(job_id, status_msg, cancel_event,
-                                               total_bytes=file_size, action="sent")
-            prog_send_task = prog_send['task']
-            send_bytes_holder = prog_send['bytes']
+            prog_task, bytes_holder = await send_progress_updater(
+                job_id, status_msg, cancel_event, file_size
+            )
 
             async def send_progress(current, total):
-                send_bytes_holder[0] = current
+                bytes_holder[0] = current
 
             await app.send_document(
                 chat_id,
@@ -363,17 +317,20 @@ async def do_download_job(chat_id: int, enc_path: str, orig_name: str,
                 caption=orig_name,
                 progress=send_progress
             )
-            prog_send_task.cancel()
+            prog_task.cancel()
             await status_msg.delete()
         else:
             part_size = MAX_FILE_SIZE_TG
             num_parts = -(-file_size // part_size)
-            await status_msg.edit_text(f"📤 Sending in {num_parts} parts...", reply_markup=None)
             with open(local_file, "rb") as f:
                 for i in range(1, num_parts + 1):
                     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}")]
+                                               ]))
                     chunk = f.read(part_size)
                     part_name = f"{orig_name}.part{i}of{num_parts}"
                     part_path = os.path.join(temp_dir, part_name)
@@ -389,25 +346,39 @@ async def do_download_job(chat_id: int, enc_path: str, orig_name: str,
     finally:
         cleanup(temp_dir)
 
-# --------------- DELETE FILE ---------------
-async def delete_user_file(chat_id: int, user_id: int, enc_path: str, token: str, query):
+# --------------- DELETE FILE HELPER ---------------
+async def delete_user_file(chat_id, user_id, enc_path, token, query):
     try:
-        # Delete from Hugging Face
         hf.delete_file(
             path_in_repo=enc_path,
             repo_id=DATASET_REPO,
             repo_type="dataset",
         )
-        # Remove from database
         db_remove_file(enc_path)
-        # Remove from callback map
         file_callback_map.pop(token, None)
-        await query.answer("File deleted successfully.", show_alert=True)
-        # Refresh the file list message
-        await list_files(_, query.message)  # we'll need to call the actual handler, but we don't have the client here; simpler: edit the message
-        # Instead, we'll just remove the button row (or update the message)
+        await query.answer("File deleted.", show_alert=True)
+        await refresh_file_list(chat_id, user_id, query.message)
     except Exception as e:
-        await query.answer(f"Error deleting file: {e}", show_alert=True)
+        await query.answer(f"Error: {e}", show_alert=True)
+
+async def refresh_file_list(chat_id, user_id, 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:
+        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}")
+        ])
+
+    await message.edit_text(
+        "Your uploaded files (click to download, delete to remove):",
+        reply_markup=InlineKeyboardMarkup(buttons)
+    )
 
 # --------------- QUEUE WORKER ---------------
 async def upload_worker():
@@ -477,12 +448,10 @@ async def list_files_cmd(_, message):
     for enc_path, orig_name, size in user_files:
         token = make_token()
         file_callback_map[token] = enc_path
-        # Two buttons: Download and Delete
-        row = [
+        buttons.append([
             InlineKeyboardButton(f"📥 {orig_name[:40]}", callback_data=f"get:{token}"),
             InlineKeyboardButton("🗑 Delete", callback_data=f"delete:{token}")
-        ]
-        buttons.append(row)
+        ])
 
     await message.reply_text(
         "Your uploaded files (click to download, delete to remove):",
@@ -532,22 +501,7 @@ async def delete_file_callback(_, query):
     user_id = query.from_user.id
     chat_id = query.message.chat.id
 
-    # Delete from HF and DB
-    try:
-        hf.delete_file(
-            path_in_repo=enc_path,
-            repo_id=DATASET_REPO,
-            repo_type="dataset",
-        )
-        db_remove_file(enc_path)
-        await query.answer("File deleted.", show_alert=True)
-
-        # Remove the button row by editing the message (optional)
-        # We'll just remove that specific button row from the keyboard.
-        # Better: refresh the whole list.
-        await list_files_cmd(_, query.message)  # this will edit the same message with the updated list
-    except Exception as e:
-        await query.answer(f"Error deleting file: {e}", show_alert=True)
+    await delete_user_file(chat_id, user_id, enc_path, token, query)
 
 @app.on_callback_query(filters.regex("^cancel_job:"))
 async def cancel_job_callback(_, query):