Update app.py

app.py

@@ -1,473 +1,234 @@
 import os
-import sys
-import time
-import base64
-import hashlib
+import tempfile
 import shutil
 import threading
-import asyncio
-import tempfile
-from pathlib import Path
-from typing import Optional
+import queue
 
-import requests
 from flask import Flask
+
 from pyrogram import Client, filters
 from pyrogram.types import InlineKeyboardMarkup, InlineKeyboardButton
-from huggingface_hub import HfApi, hf_hub_download, create_commit, CommitOperationAdd
-from cryptography.fernet import Fernet
 
-# --------------- CONFIG ---------------
+from huggingface_hub import HfApi, hf_hub_download
+
+# =========================
+# 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")
 
 DATASET_REPO = "Cxops/Bzxv"
-MAX_FILE_SIZE_TG = 1900 * 1024 * 1024   # 1900 MB per part (safe margin)
-MIN_FREE_DISK = 45 * 1024 * 1024 * 1024  # 45 GB
-KEEPALIVE_INTERVAL = 300                 # 5 minutes
-
-# --------------- ENCRYPTION SETUP ---------------
-fernet_key = base64.urlsafe_b64encode(hashlib.sha256(OBFUSCATION_KEY.encode()).digest())
-fernet = Fernet(fernet_key)
-
-def encrypt_text(plain: str) -> str:
-    token = fernet.encrypt(plain.encode())
-    # URL‑safe base64 without padding (safe for filenames)
-    return base64.urlsafe_b64encode(token).decode().rstrip("=")
-
-def decrypt_text(encoded: str) -> str:
-    # Add padding back
-    padding = 4 - len(encoded) % 4
-    if padding != 4:
-        encoded += "=" * padding
-    token = base64.urlsafe_b64decode(encoded.encode())
-    return fernet.decrypt(token).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)
-
-# --------------- FLASK HEALTH SERVER ---------------
+
+# =========================
+# FLASK (HEALTH CHECK)
+# =========================
+
 web = Flask(__name__)
 
 @web.route("/")
 def home():
-    return "Telegram HF Storage Bot Running"
+    return "OK - Ultra Fast Bot Running"
 
 def run_web():
     web.run(host="0.0.0.0", port=7860)
 
-# --------------- HF CLIENT ---------------
+# =========================
+# HF CLIENT
+# =========================
+
 hf = HfApi(token=HF_TOKEN)
 
-# --------------- PYROGRAM BOT ---------------
+# =========================
+# UPLOAD QUEUE SYSTEM (FAST)
+# =========================
+
+upload_queue = queue.Queue()
+
+def worker():
+    while True:
+        try:
+            local_path, repo_path, status = upload_queue.get()
+
+            hf.upload_file(
+                path_or_fileobj=local_path,
+                path_in_repo=repo_path,
+                repo_id=DATASET_REPO,
+                repo_type="dataset",
+            )
+
+            try:
+                status.edit_text("⚡ Upload Complete (Ultra Fast)")
+            except:
+                pass
+
+        except Exception as e:
+            try:
+                status.edit_text(f"Upload Failed:\n{e}")
+            except:
+                pass
+
+        finally:
+            upload_queue.task_done()
+
+# start multiple workers (PARALLEL SPEED)
+for _ in range(3):
+    threading.Thread(target=worker, daemon=True).start()
+
+# =========================
+# PYROGRAM
+# =========================
+
 app = Client(
-    "storage_bot",
+    "ultra_fast_bot",
     api_id=API_ID,
     api_hash=API_HASH,
     bot_token=BOT_TOKEN
 )
 
-# --------------- ASYNC QUEUES & STATE ---------------
-upload_queue = asyncio.Queue()
-active_jobs = {}          # job_id -> {status, cancel_event, message, ...}
-processing_now = False    # simple mutex for the worker
+# =========================
+# CLEANUP
+# =========================
 
-# --------------- HELPERS ---------------
-
-def cleanup(path: str):
+def cleanup(path):
     try:
-        if os.path.isdir(path):
-            shutil.rmtree(path, ignore_errors=True)
-        elif os.path.isfile(path):
-            os.remove(path)
-    except Exception:
+        shutil.rmtree(path, ignore_errors=True)
+    except:
         pass
 
-def get_free_disk() -> int:
-    """Return free space in bytes on the current filesystem."""
-    st = shutil.disk_usage("/")
-    return st.free
-
-async def progress_updater(job_id: int, msg, cancel_event: asyncio.Event, total_bytes: Optional[int] = None):
-    """Update message every 5 seconds with a progress indicator (elapsed time or percentage)."""
-    start = time.time()
-    while not cancel_event.is_set():
-        elapsed = int(time.time() - start)
-        text = f"⏳ Working... {elapsed}s elapsed"
-        if total_bytes:
-            # can't show real bytes, just a generic message
-            pass
-        try:
-            await msg.edit_text(text, reply_markup=InlineKeyboardMarkup([
-                [InlineKeyboardButton("❌ Cancel", callback_data=f"cancel_job:{job_id}")]
-            ]))
-        except Exception:
-            break
-        await asyncio.sleep(5)
-    # Final cancellation message will be set by the caller
-
-# --------------- UPLOAD / DOWNLOAD LOGIC ---------------
-
-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):
-    """Download from Telegram, upload to HF, update progress, handle cancel."""
-    temp_dir = tempfile.mkdtemp()
-    try:
-        # 1. Download file from Telegram
-        await status_msg.edit_text("⬇️ Downloading from Telegram...",
-                                   reply_markup=InlineKeyboardMarkup([
-                                       [InlineKeyboardButton("❌ Cancel", callback_data=f"cancel_job:{job_id}")]
-                                   ]))
-        local_path = await file_msg.download(file_name=os.path.join(temp_dir, orig_filename))
-        if cancel_event.is_set():
-            return
-
-        # 2. Upload to Hugging Face dataset
-        await status_msg.edit_text("⬆️ Uploading to Hugging Face...",
-                                   reply_markup=InlineKeyboardMarkup([
-                                       [InlineKeyboardButton("❌ Cancel", callback_data=f"cancel_job:{job_id}")]
-                                   ]))
-        # Encrypt folder + filename
-        enc_folder = get_encrypted_folder(user_id)
-        enc_filename = get_encrypted_filename(orig_filename)
-        repo_path = f"{enc_folder}/{enc_filename}"
-
-        # Use a synchronous upload in an executor, keep updating progress
-        loop = asyncio.get_event_loop()
-        upload_done = loop.create_future()
-
-        def sync_upload():
-            try:
-                hf.upload_file(
-                    path_or_fileobj=local_path,
-                    path_in_repo=repo_path,
-                    repo_id=DATASET_REPO,
-                    repo_type="dataset",
-                )
-                return True
-            except Exception as e:
-                raise e
-
-        task = loop.run_in_executor(None, sync_upload)
-
-        # Progress task runs in background
-        progress_task = asyncio.create_task(progress_updater(job_id, status_msg, cancel_event))
-
-        # Wait for upload to finish or cancel
-        while not task.done():
-            if cancel_event.is_set():
-                # Cancel is set – we cannot stop the in‑flight upload, but we can delete the file afterwards
-                break
-            await asyncio.sleep(1)
-
-        # If cancelled while uploading, we still wait for the upload to finish (or timeout) so we can delete it
-        if cancel_event.is_set():
-            # Wait for upload to finish (with a timeout)
-            try:
-                await asyncio.wait_for(task, timeout=30)   # hope it finishes soon
-            except asyncio.TimeoutError:
-                pass  # can't delete if still running, but at least we tried
-            # Delete the uploaded file from HF
-            try:
-                hf.delete_file(
-                    path_in_repo=repo_path,
-                    repo_id=DATASET_REPO,
-                    repo_type="dataset",
-                )
-            except Exception:
-                pass
-            progress_task.cancel()
-            await status_msg.edit_text("❌ Upload cancelled.",
-                                       reply_markup=None)
-            return
+# =========================
+# START
+# =========================
 
-        # Normal completion
-        await task  # raise any exception
-        progress_task.cancel()
-        await status_msg.edit_text(f"✅ Uploaded: `{orig_filename}`",
-                                   reply_markup=None)
+@app.on_message(filters.command("start"))
+async def start(_, m):
+    await m.reply_text(
+        "⚡ ULTRA FAST STORAGE BOT\n\n"
+        "Send file → stored in Hugging Face\n"
+        "/files → list files"
+    )
 
-    except Exception as e:
-        await status_msg.edit_text(f"❌ Upload failed: {e}",
-                                   reply_markup=None)
-    finally:
-        cleanup(temp_dir)
+# =========================
+# LIST FILES
+# =========================
+
+@app.on_message(filters.command("files"))
+async def files(_, m):
+    uid = str(m.from_user.id)
 
-async def do_download_job(chat_id: int, enc_path: str, job_id: int, cancel_event: asyncio.Event, status_msg):
-    """Download from HF, split if >2GB, send parts to user, progress updates."""
-    temp_dir = tempfile.mkdtemp()
     try:
-        # Decrypt original filename from path
-        _, enc_filename = enc_path.split("/", 1)
-        orig_filename = decrypt_text(enc_filename)
+        all_files = hf.list_repo_files(
+            repo_id=DATASET_REPO,
+            repo_type="dataset"
+        )
 
-        await status_msg.edit_text("⬇️ Downloading from dataset...",
-                                   reply_markup=InlineKeyboardMarkup([
-                                       [InlineKeyboardButton("❌ Cancel", callback_data=f"cancel_job:{job_id}")]
-                                   ]))
+        user_files = [f for f in all_files if f.startswith(uid + "/")]
 
-        loop = asyncio.get_event_loop()
-        download_done = loop.create_future()
+        if not user_files:
+            return await m.reply_text("No files found.")
 
-        def sync_download():
-            return hf_hub_download(
-                repo_id=DATASET_REPO,
-                repo_type="dataset",
-                filename=enc_path,
-                token=HF_TOKEN,
-                local_dir=temp_dir,
-            )
+        buttons = [
+            [InlineKeyboardButton(os.path.basename(f), callback_data=f"get:{f}")]
+            for f in user_files[:50]
+        ]
 
-        task = loop.run_in_executor(None, sync_download)
-        progress_task = asyncio.create_task(progress_updater(job_id, status_msg, cancel_event))
-
-        while not task.done():
-            if cancel_event.is_set():
-                # wait a bit then cancel
-                break
-            await asyncio.sleep(1)
-
-        if cancel_event.is_set():
-            progress_task.cancel()
-            await status_msg.edit_text("❌ Download cancelled.",
-                                       reply_markup=None)
-            return
-
-        local_file = await task
-        progress_task.cancel()
-
-        # Check file size
-        file_size = os.path.getsize(local_file)
-        if file_size <= MAX_FILE_SIZE_TG:
-            # Send as single file
-            await status_msg.edit_text("📤 Sending file to Telegram...",
-                                       reply_markup=None)
-            await app.send_document(
-                chat_id,
-                document=local_file,
-                file_name=orig_filename,
-                caption=orig_filename
-            )
-            await status_msg.delete()
-        else:
-            # Split into parts
-            await status_msg.edit_text("📤 Splitting and sending...",
-                                       reply_markup=None)
-            part_size = MAX_FILE_SIZE_TG
-            num_parts = -(-file_size // part_size)  # ceil division
-            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
-                    chunk = f.read(part_size)
-                    part_name = f"{orig_filename}.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_filename} – Part {i}/{num_parts}"
-                    )
-                    os.remove(part_path)
-            await status_msg.delete()
+        await m.reply_text(
+            "Your files:",
+            reply_markup=InlineKeyboardMarkup(buttons)
+        )
 
     except Exception as e:
-        await status_msg.edit_text(f"❌ Download failed: {e}")
-    finally:
-        cleanup(temp_dir)
+        await m.reply_text(str(e))
 
-# --------------- QUEUE WORKER ---------------
-async def upload_worker():
-    """Process upload jobs one by one, respecting disk space."""
-    global processing_now
-    while True:
-        job = await upload_queue.get()
-        user_id, chat_id, file_msg, orig_filename, job_id = job
-        cancel_event = asyncio.Event()
-        active_jobs[job_id]["cancel_event"] = cancel_event
-        active_jobs[job_id]["status"] = "running"
-
-        # Check disk space
-        while get_free_disk() < MIN_FREE_DISK:
-            # Not enough space, wait a bit
-            await asyncio.sleep(10)
-
-        # Create a new status message
-        status_msg = await app.send_message(
-            chat_id,
-            "⏳ Starting upload...",
-            reply_markup=InlineKeyboardMarkup([
-                [InlineKeyboardButton("❌ Cancel", callback_data=f"cancel_job:{job_id}")]
-            ])
+# =========================
+# DOWNLOAD FROM HF
+# =========================
+
+@app.on_callback_query(filters.regex("^get:"))
+async def get_file(_, q):
+    path = q.data.split(":", 1)[1]
+
+    tmp = tempfile.mkdtemp()
+
+    try:
+        await q.message.edit_text("⚡ Downloading...")
+
+        file_path = hf_hub_download(
+            repo_id=DATASET_REPO,
+            repo_type="dataset",
+            filename=path,
+            token=HF_TOKEN,
+            local_dir=tmp
         )
-        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)
-        finally:
-            # Cleanup
-            active_jobs.pop(job_id, None)
-            upload_queue.task_done()
+        await q.message.reply_document(file_path)
 
-# --------------- BOT HANDLERS ---------------
+        await q.message.delete()
 
-@app.on_message(filters.command("start"))
-async def start_cmd(_, message):
-    await message.reply_text(
-        "Send me a file to upload it to the dataset.\n"
-        "/files – list your uploaded files"
-    )
+    except Exception as e:
+        await q.message.reply_text(str(e))
 
-@app.on_message(filters.command("files"))
-async def list_files(_, message):
-    user_id = message.from_user.id
-    enc_folder = get_encrypted_folder(user_id)
+    finally:
+        cleanup(tmp)
 
-    try:
-        all_files = hf.list_repo_files(repo_id=DATASET_REPO, repo_type="dataset")
-        user_files = [f for f in all_files if f.startswith(enc_folder + "/")]
-        if not user_files:
-            return await message.reply_text("You have no uploaded files.")
+# =========================
+# UPLOAD FROM TELEGRAM (FAST)
+# =========================
 
-        buttons = []
-        for path in user_files:
-            _, enc_name = path.split("/", 1)
-            try:
-                orig_name = decrypt_text(enc_name)
-            except Exception:
-                orig_name = enc_name  # fallback
-            buttons.append([InlineKeyboardButton(orig_name[:60], callback_data=f"get:{path}")])
+@app.on_message(
+    filters.document |
+    filters.video |
+    filters.audio |
+    filters.photo
+)
+async def upload(_, m):
 
-        await message.reply_text(
-            "Your uploaded files:",
-            reply_markup=InlineKeyboardMarkup(buttons)
-        )
-    except Exception as e:
-        await message.reply_text(f"Error listing files: {e}")
+    tmp = tempfile.mkdtemp()
 
-@app.on_callback_query(filters.regex("^get:"))
-async def get_file_callback(_, query):
-    path = query.data.replace("get:", "")
-    job_id = int(time.time() * 1000)  # simple unique id
-    cancel_event = asyncio.Event()
-    active_jobs[job_id] = {
-        "status": "running",
-        "cancel_event": cancel_event,
-        "type": "download"
-    }
-
-    status_msg = await query.message.edit_text(
-        "⏳ Preparing download...",
-        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,
-        path,
-        job_id,
-        cancel_event,
-        status_msg
-    ))
-
-@app.on_callback_query(filters.regex("^cancel_job:"))
-async def cancel_job_callback(_, query):
-    _, job_id_str = query.data.split(":")
-    job_id = int(job_id_str)
-    if job_id in active_jobs:
-        active_jobs[job_id]["cancel_event"].set()
-        await query.answer("Cancelling...")
-    else:
-        await query.answer("Job not found.")
-
-@app.on_callback_query(filters.regex("^upload:"))
-async def upload_callback(_, query):
-    # Data: upload:user_id:message_id:orig_filename
-    _, user_id_str, msg_id_str, orig_filename = query.data.split(":", 3)
-    user_id = int(user_id_str)
-    chat_id = query.message.chat.id
-
-    # Find the original file message
     try:
-        file_msg = await app.get_messages(chat_id, int(msg_id_str))
-    except Exception:
-        return await query.answer("Original message missing.", show_alert=True)
+        status = await m.reply_text("⚡ Downloading...")
 
-    if not file_msg:
-        return await query.answer("File no longer available.", show_alert=True)
+        if m.document:
+            name = m.document.file_name
 
-    job_id = int(time.time() * 1000)
-    active_jobs[job_id] = {"status": "queued"}
+        elif m.video:
+            name = m.video.file_name or f"{m.video.file_unique_id}.mp4"
 
-    await upload_queue.put((user_id, chat_id, file_msg, orig_filename, job_id))
+        elif m.audio:
+            name = m.audio.file_name or f"{m.audio.file_unique_id}.mp3"
 
-    await query.message.edit_text(
-        f"📦 `{orig_filename}` added to upload queue.\n"
-        "You will be notified when it starts.",
-        reply_markup=None
-    )
+        else:
+            name = f"{m.photo[-1].file_unique_id}.jpg"
 
-@app.on_callback_query(filters.regex("^cancel_upload:"))
-async def cancel_upload_prompt(_, query):
-    await query.message.delete()
-
-@app.on_message(filters.document | filters.video | filters.audio | filters.photo)
-async def on_file(_, message):
-    # Determine file name
-    if message.document:
-        filename = message.document.file_name
-    elif message.video:
-        filename = message.video.file_name or f"{message.video.file_unique_id}.mp4"
-    elif message.audio:
-        filename = message.audio.file_name or f"{message.audio.file_unique_id}.mp3"
-    else:  # photo
-        filename = f"{message.photo.file_unique_id}.jpg"
-
-    # Build inline keyboard
-    kb = InlineKeyboardMarkup([
-        [InlineKeyboardButton("⬆️ Upload File", callback_data=f"upload:{message.from_user.id}:{message.id}:{filename}")],
-        [InlineKeyboardButton("❌ Cancel", callback_data="cancel_upload:0")]
-    ])
-
-    await message.reply_text(
-        f"Do you want to upload `{filename}`?",
-        reply_markup=kb
-    )
+        local_path = await m.download(file_name=f"{tmp}/{name}")
 
-# --------------- KEEPALIVE ---------------
-def keepalive():
-    while True:
-        try:
-            requests.get("http://localhost:7860/", timeout=10)
-        except Exception:
-            pass
-        time.sleep(KEEPALIVE_INTERVAL)
+        repo_path = f"{m.from_user.id}/{name}"
+
+        # 🔥 push into queue (INSTANT RETURN)
+        upload_queue.put((local_path, repo_path, status))
+
+        await status.edit_text("⚡ Queued for ultra-fast upload...")
 
-# --------------- MAIN ---------------
-async def main():
-    # Start the upload worker
-    asyncio.create_task(upload_worker())
+    except Exception as e:
+        await m.reply_text(f"Error:\n{e}")
+
+    finally:
+        def delayed_cleanup():
+            import time
+            time.sleep(15)
+            cleanup(tmp)
 
-    # Start the bot
-    await app.start()
-    print("Bot started")
-    await asyncio.Event().wait()  # run forever
+        threading.Thread(target=delayed_cleanup, daemon=True).start()
+
+# =========================
+# RUN
+# =========================
 
 if __name__ == "__main__":
-    # Start Flask in a separate thread
+    print("🚀 Ultra Fast Bot Starting...")
+
     threading.Thread(target=run_web, daemon=True).start()
-    # Start keepalive pinger
-    threading.Thread(target=keepalive, daemon=True).start()
 
-    # Run Pyrogram
-    app.run(main())
+    app.run()