main.py

import os
import time
import uuid
import asyncio
import threading
import logging
from flask import Flask, redirect
from pyrogram import Client, filters, idle
from pyrogram.types import InlineKeyboardMarkup, InlineKeyboardButton
from huggingface_hub import HfApi

# --- LOGGING ---
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("Rajasthan_Bot")

# --- ENV VARIABLES ---
def get_clean_var(name):
    val = os.environ.get(name, "")
    return val.strip() if val else ""

API_ID = get_clean_var("API_ID")
API_HASH = get_clean_var("API_HASH")
BOT_TOKEN = get_clean_var("BOT_TOKEN")
SESSION_STRING = get_clean_var("SESSION_STRING") 
HF_TOKEN = get_clean_var("HF_TOKEN") 
HF_REPO = get_clean_var("HF_REPO")   
ACCESS_PASSWORD = get_clean_var("PASSWORD") or "Maharaja Jaswant Singh"

SPACE_HOST = os.environ.get("SPACE_HOST", "localhost:7860")
BASE_URL = f"https://{SPACE_HOST}"

# --- FLASK SERVER ---
app = Flask(__name__)

@app.route('/')
def home():
    return "Rajasthan Bot System Online 🟢"

@app.route('/rajasthan/<path:filename>')
def serve_file(filename):
    real_link = f"https://huggingface.co/datasets/{HF_REPO}/resolve/main/{filename}?download=true"
    return redirect(real_link, code=302)

def run_web_server():
    app.run(host="0.0.0.0", port=7860)

# --- HELPERS ---
def humanbytes(size):
    if not size: return "0 B"
    power = 2**10
    n = 0
    power_labels = {0 : '', 1: 'KB', 2: 'MB', 3: 'GB', 4: 'TB'}
    while size > power:
        size /= power
        n += 1
    return f"{size:.2f} {power_labels[n]}"

# --- GLOBAL VARS ---
process_lock = asyncio.Lock()
AUTH_USERS = set()
USERBOT_ALIVE = False

# --- CLIENTS ---
if not API_ID or not BOT_TOKEN:
    print("❌ ERROR: API_ID ya BOT_TOKEN missing hai!")
    exit(1)

try:
    API_ID = int(API_ID)
except:
    print("❌ ERROR: API_ID number hona chahiye.")
    exit(1)

bot = Client("main_bot", api_id=API_ID, api_hash=API_HASH, bot_token=BOT_TOKEN)

userbot = None
if SESSION_STRING:
    userbot = Client("user_bot", api_id=API_ID, api_hash=API_HASH, session_string=SESSION_STRING)

# --- PROGRESS ---
async def progress(current, total, message, start_time, status_text):
    now = time.time()
    diff = now - start_time
    if round(diff % 5.00) == 0 or current == total:
        percentage = current * 100 / total
        speed = current / diff if diff > 0 else 0
        try:
            await message.edit(
                f"{status_text}\n"
                f"━━━━━━━━━━━━━━━━━━\n"
                f"📊 **Progress:** {percentage:.1f}%\n"
                f"💾 **Done:** {humanbytes(current)} / {humanbytes(total)}\n"
                f"⚡ **Speed:** {humanbytes(speed)}/s"
            )
        except:
            pass

# --- PROCESS LOGIC ---
async def process_media(client, media_msg, status_msg, user_request_msg):
    unique_id = uuid.uuid4().hex[:5]
    
    media = media_msg.video or media_msg.audio or media_msg.photo or media_msg.document
    file_size_bytes = getattr(media, "file_size", 0)
    readable_size = humanbytes(file_size_bytes)

    if media_msg.video:
        ext = "mp4"
        name_type = "Video"
    elif media_msg.audio:
        ext = "mp3"
        name_type = "Music"
    elif media_msg.photo:
        ext = "jpg"
        name_type = "Image"
    elif media_msg.document:
        try: ext = media_msg.document.file_name.split(".")[-1]
        except: ext = "pdf"
        name_type = "File"
    else:
        ext = "file"
        name_type = "File"

    filename = f"Rajasthan_{name_type}_{unique_id}.{ext}"
    save_path = f"./{filename}"

    try:
        start = time.time()
        await status_msg.edit(f"⬇️ **Downloading...**\n`{filename}`")
        
        await client.download_media(
            message=media_msg,
            file_name=save_path,
            progress=progress,
            progress_args=(status_msg, start, "⬇️ **Downloading (Userbot)...**")
        )

        await status_msg.edit("☁️ **Uploading to Cloud...**")
        api = HfApi(token=HF_TOKEN)
        await asyncio.to_thread(
            api.upload_file,
            path_or_fileobj=save_path,
            path_in_repo=filename,
            repo_id=HF_REPO,
            repo_type="dataset"
        )

        final_link = f"{BASE_URL}/rajasthan/{filename}"
        
        await status_msg.delete()
        
        hyperlink_text = f"[{filename}]({final_link})"

        await user_request_msg.reply_text(
            f"⚡ **GENERATED SUCCESSFULLY**\n"
            f"━━━━━━━━━━━━━━━━━━\n"
            f"📂 **File:** {hyperlink_text}\n"
            f"💾 **Size:** `{readable_size}`\n"
            f"━━━━━━━━━━━━━━━━━━\n"
            f"🔗 **Link:** `{final_link}`",
            disable_web_page_preview=True,
            reply_markup=InlineKeyboardMarkup([
                [InlineKeyboardButton("🚀 One Click Download", url=final_link)]
            ])
        )

    except Exception as e:
        logger.error(f"Error: {e}")
        await status_msg.edit(f"❌ **Error:** {str(e)}")
    
    finally:
        if os.path.exists(save_path): os.remove(save_path)

# --- HANDLERS ---
@bot.on_message(filters.command("start"))
async def start(c, m):
    if m.from_user.id in AUTH_USERS:
        await m.reply_text("👋 **Welcome Back!**")
    else:
        await m.reply_text("🔒 **Access Denied!** Enter Password.")

@bot.on_message(filters.private & filters.text)
async def text_handler(c, m):
    if m.from_user.id not in AUTH_USERS:
        if m.text == ACCESS_PASSWORD:
            AUTH_USERS.add(m.from_user.id)
            await m.reply_text("🔓 **Access Granted!**")
        else:
            await m.reply_text("❌ **Wrong Password!**")
        return

    if "t.me/" in m.text:
        if not userbot or not USERBOT_ALIVE: return await m.reply_text("⚠️ **Userbot Error.**")
        if process_lock.locked(): return await m.reply_text("⚠️ **Queue Full.**")

        async with process_lock:
            status = await m.reply_text("🔎 **Deep Scanning (Range: 100)...**")
            try:
                # 1. Clean Link
                link = m.text.strip().replace("https://", "").replace("http://", "")
                if "t.me/" in link: link = link.split("t.me/")[1] 
                
                parts = link.split("/")
                
                if parts[0] == "c":
                    chat_id = int("-100" + parts[1]) # Private
                else:
                    chat_id = parts[0] # Public Username

                start_msg_id = int(parts[-1].split("?")[0])

                print(f"DEBUG: Chat: {chat_id}, Start ID: {start_msg_id}")

                # 2. DEEP SCAN: Check Start ID + Next 100 messages
                # Topic ID gaps can be huge in busy groups
                
                target_msg = None
                found_at_id = 0
                check_limit = 100 # Increased from 10 to 100
                
                try:
                    messages = await userbot.get_messages(chat_id, range(start_msg_id, start_msg_id + check_limit))
                except Exception as e:
                    return await status.edit(f"❌ **Scan Error:** {e}")

                if not isinstance(messages, list):
                    messages = [messages]

                for msg in messages:
                    if msg:
                        # Check 1: Direct Media
                        if msg.media and not msg.web_page:
                            target_msg = msg
                            found_at_id = msg.id
                            break
                        
                        # Check 2: If message is a Reply to a file
                        if msg.reply_to_message and msg.reply_to_message.media:
                            target_msg = msg.reply_to_message
                            found_at_id = msg.reply_to_message.id
                            break

                if not target_msg:
                    return await status.edit(
                        f"❌ **No Media Found in Range!**\n"
                        f"Bot scanned from ID `{start_msg_id}` to `{start_msg_id + check_limit}`.\n\n"
                        f"**Reason:** In Topic groups, Message IDs are shared across all topics. The file might be very far down.\n"
                        f"👉 **Try:** Copy the link of the File itself, NOT the heading."
                    )

                if found_at_id != start_msg_id:
                     await status.edit(f"✅ **Found File at ID:** `{found_at_id}`\n(Scanned forward due to Topic gaps)\n\n⬇️ **Processing...**")
                
                await process_media(userbot, target_msg, status, m)
                
            except Exception as e: 
                await status.edit(f"❌ **Error:** {e}")

@bot.on_message(filters.private & (filters.document | filters.video | filters.audio | filters.photo))
async def file_handler(c, m):
    if m.from_user.id not in AUTH_USERS: return await m.reply_text("🔒 **Password Required.**")
    if process_lock.locked(): return await m.reply_text("⚠️ **Queue Full.**")

    async with process_lock:
        status = await m.reply_text("⏳ **Added to Queue...**")
        await process_media(bot, m, status, m)

# --- STARTUP ---
async def main():
    threading.Thread(target=run_web_server, daemon=True).start()
    
    print("🚀 Bot Starting...")
    try:
        await bot.start()
        print("✅ Main Bot Connected Successfully!")
    except Exception as e:
        print(f"❌ Main Bot Start Error: {e}")
        return

    global USERBOT_ALIVE
    if userbot:
        try:
            await userbot.start()
            USERBOT_ALIVE = True
            print("✅ Userbot Connected!")
        except Exception as e:
            print(f"⚠️ Userbot Failed: {e}")
            USERBOT_ALIVE = False
            
    await idle()
    await bot.stop()

if __name__ == "__main__":
    loop = asyncio.get_event_loop()
    loop.run_until_complete(main())