# --- START OF FILE main.py ---

import os
import time
import uuid
import asyncio
import threading
import logging
import hashlib
import re
import requests
from flask import Flask
from flask_cors import CORS
from pyrogram import Client, filters, idle
from pyrogram.errors import FloodWait  # 🔥 NEW: Anti-Flood system import
from internetarchive import upload

# --- 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") 
ACCESS_PASSWORD = get_clean_var("PASSWORD") or "Maharaja Jaswant Singh"

IA_ACCESS_KEY = get_clean_var("IA_ACCESS_KEY")
IA_SECRET_KEY = get_clean_var("IA_SECRET_KEY")
IA_BUCKET = get_clean_var("IA_BUCKET")

# 🔥 FIREBASE URL FIX 🔥
FIREBASE_URL = get_clean_var("FIREBASE_URL")
if FIREBASE_URL and not FIREBASE_URL.startswith("http"):
    FIREBASE_URL = "https://" + FIREBASE_URL
if FIREBASE_URL.endswith("/"):
    FIREBASE_URL = FIREBASE_URL[:-1] 

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

@app.route('/')
def home():
    return "Admin Pro - Auto Sync Bot Online 🟢"

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

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

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

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

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

# --- 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]}"

def get_id(name):
    return hashlib.md5(name.strip().lower().encode()).hexdigest()[:12]

def get_base_title(title):
    t = str(title).lower()
    t = re.sub(r'\(pdf.*?\)', '', t)
    t = re.sub(r'\(notes\)', '', t)
    t = re.sub(r'pdf\s*\d*', '', t)
    t = re.sub(r'part\s*\d+', '', t)
    t = re.sub(r'\s*\d+\s*$', '', t) 
    t = re.sub(r'[^\w]', '', t)  
    return t.strip()

# 🔥 CAPTION PARSER 🔥
def parse_caption(caption, is_pdf=False):
    batch = "Unknown Batch"
    subject = "General Subject"
    chapter = "General Chapter"
    title = "Class Lecture"

    if not caption:
        return batch, subject, chapter, title

    try:
        # 1. Batch Nikalna
        batch_match = re.search(r'Batch:\s*([^\n]+)', caption, re.IGNORECASE)
        if batch_match: 
            batch = batch_match.group(1).strip()

        # 2. Subject & Chapter (Agar direct Topic: likha ho)
        topic_match = re.search(r'Topic:\s*([^\n]+)', caption, re.IGNORECASE)
        if topic_match:
            raw_topic = topic_match.group(1).strip()
            if '→' in raw_topic: parts = raw_topic.split('→')
            elif '->' in raw_topic: parts = raw_topic.split('->')
            elif '-' in raw_topic: parts = raw_topic.split('-')
            else: parts = [raw_topic]

            if len(parts) >= 2:
                subject = parts[0].strip()
                chapter = re.sub(r'[•·.]+$', '', parts[1].strip()).strip()
            else:
                subject = parts[0].strip()

        # 3. Title & HINDI CHAPTER Logic
        raw_title = ""
        title_line_match = re.search(r'Title:\s*([^\n]+)', caption, re.IGNORECASE)
        
        if title_line_match:
            raw_title = title_line_match.group(1).strip()
        else:
            # Agar "Title:" nahi likha hai, to pehli line ko title maan lo
            lines = [line.strip() for line in caption.split('\n') if line.strip()]
            if lines: raw_title = lines[0]

        if raw_title:
            base_title = re.sub(r'\.mkv|\.mp4|\.pdf', '', raw_title, flags=re.IGNORECASE).strip()

            # 🔥 NEW MAGIC: Hindi words ko scan karke Chapter banana 🔥
            # \u0900-\u097F ye code computer ko Hindi padhna sikhata hai
            hindi_matches = re.findall(r'[\u0900-\u097F]+(?:\s+[\u0900-\u097F]+)*', base_title)
            
            if hindi_matches and chapter == "General Chapter":
                # Agar title me hindi word mila (jaise: 'संविधान सभा'), to usko Chapter set kardo
                chapter = max(hindi_matches, key=len).strip()

            if is_pdf:
                if "pdf 2" in raw_title.lower(): title = f"{base_title} (PDF 2)"
                elif "pdf" in raw_title.lower(): title = f"{base_title} (PDF 1)"
                else: title = f"{base_title} (Notes)"
            else:
                title = base_title

    except Exception as e:
        logger.error(f"Caption parsing error: {e}")

    return batch, subject, chapter, title

# 🔥 PROGRESS TRACKER (WITH ANTI-FLOOD FIX) 🔥
async def progress(current, total, message, start_time, header_card):
    now = time.time()
    diff = now - start_time
    # Telegram spam limits ke liye update time ko 5 seconds kar diya gaya hai
    if round(diff % 5.0) == 0 or current == total:
        percentage = current * 100 / total
        speed = current / diff if diff > 0 else 0
        text = (
            f"{header_card}\n"
            f"━━━━━━━━━━━━━━━━━━\n"
            f"⬇️ **Status:** Downloading from Telegram...\n"
            f"📊 **Progress:** {percentage:.1f}%\n"
            f"💾 **Done:** {humanbytes(current)} / {humanbytes(total)}\n"
            f"⚡ **Speed:** {humanbytes(speed)}/s"
        )
        try:
            await message.edit(text)
        except FloodWait as e:
            # Agar telegram rokega, toh bot shaanti se wait karega!
            await asyncio.sleep(e.value)
        except Exception:
            pass

# --- COMMANDS ---
@bot.on_message(filters.command("start"))
async def start_cmd(c, m):
    if m.from_user.id in AUTH_USERS:
        await m.reply_text("👋 **Welcome Back!**\nDirect Video/PDF ki link bhejo (Jaise msg ID 866 waali).")
    else:
        await m.reply_text("🔒 **Access Denied!** Pehle Password bhejo.")

@bot.on_message(filters.command("stop"))
async def stop_cmd(c, m):
    global STOP_REQUESTED
    if m.from_user.id not in AUTH_USERS: return
    STOP_REQUESTED = True
    await m.reply_text("🛑 **Stop Command Received!** Current upload ke baad loop ruk jayega.")

# --- FIREBASE SECURE PATCHER ---
def patch_firebase(url, payload):
    res = requests.patch(url, json=payload)
    if res.status_code >= 400:
        raise Exception(f"Firebase Rejected Request (Status {res.status_code}): {res.text}")

# 🔥 DYNAMIC FOLDER RESOLVER (Chronological Sorting for CHAPTERS only) 🔥
global_firebase_cache = {}

def fetch_firebase_keys(parent_url):
    if parent_url in global_firebase_cache:
        return
    global_firebase_cache[parent_url] = {}
    if not FIREBASE_URL: return
    try:
        res = requests.get(parent_url + ".json")
        if res.status_code == 200 and res.json():
            for k, v in res.json().items():
                if isinstance(v, dict) and 'title' in v:
                    global_firebase_cache[parent_url][v['title']] = k
    except Exception as e:
        pass

def resolve_fb_id(parent_url, title, msg_id):
    fetch_firebase_keys(parent_url)
    if title in global_firebase_cache[parent_url]:
        return global_firebase_cache[parent_url][title]
    else:
        new_id = f"z_{msg_id:08d}_{get_id(title)[:6]}"
        global_firebase_cache[parent_url][title] = new_id
        return new_id

# --- MAIN LOGIC ---
@bot.on_message(filters.private & filters.text)
async def text_handler(c, m):
    global STOP_REQUESTED
    
    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!** Ab starting video ki link bhejo.")
        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 Offline hai.** Naya Session String Check Karein.")
        if process_lock.locked():
            return await m.reply_text("⚠️ **Ek extraction already chal raha hai.** Pehle `/stop` bhejo.")

        async with process_lock:
            STOP_REQUESTED = False
            status = await m.reply_text("🔎 **Admin Pro Engine Started...**")

            raw_link = m.text.strip().replace("https://", "").replace("http://", "")
            if "t.me/" in raw_link: raw_link = raw_link.split("t.me/")[1]
            parts = raw_link.split("/")

            try:
                parts = [p for p in parts if p]
                
                # 🔥 MASTER FIX: Link chahe jaisi ho, aakhri part hamesha message ID hota hai
                current_msg_id = int(parts[-1].split("?")[0])
                topic_or_post_id = None

                if parts[0] == "c":
                    # Private/Restricted Link ke liye
                    chat_id = int("-100" + parts[1])
                else:
                    # Public Link ke liye (yakeenbatchrwa jaise)
                    chat_id = parts[0]
                    
            except Exception:
                return await status.edit("❌ **Invalid Telegram Link Format!**")

            no_media_count = 0
            extracted_count = 0
            video_memory_cache = {} 
            last_fallback_video = None
            
            stop_keywords = ["chapter complete", "stop extract", "subject complete", "topic complete"]

            while True:
                if STOP_REQUESTED:
                    await m.reply_text(f"🛑 **Extraction Stopped!**\nUploaded: `{extracted_count}`")
                    break

                try:
                    msg_ids_to_check = list(range(current_msg_id, current_msg_id + 5))
                    messages = await userbot.get_messages(chat_id, msg_ids_to_check)
                    
                    if messages and messages[0].empty and no_media_count == 0 and len(parts) >= 4 and topic_or_post_id:
                        await status.edit(f"⚠️ **Message #{current_msg_id} khali hai.**\nMain-Post #{topic_or_post_id} check kar raha hu...")
                        current_msg_id = topic_or_post_id
                        msg_ids_to_check = list(range(current_msg_id, current_msg_id + 5))
                        messages = await userbot.get_messages(chat_id, msg_ids_to_check)
                        topic_or_post_id = None

                except FloodWait as e:
                    await asyncio.sleep(e.value)
                    messages = await userbot.get_messages(chat_id, msg_ids_to_check)
                except Exception as e:
                    err_msg = str(e).lower()
                    if "peer_id_invalid" in err_msg or "peer id invalid" in err_msg:
                        await status.edit("🔄 **Peer ID Error Aaya!**\nRuko, Bot apni memory (Chat Cache) auto-refresh kar raha hai...")
                        try:
                            # 🔥 FIX: Userbot automatically apni chat list fetch karke group ko pehchan lega
                            async for dialog in userbot.get_dialogs(limit=150):
                                if dialog.chat.id == chat_id:
                                    break
                            await status.edit("✅ **Cache Refresh Ho Gaya!** Ab extraction start ho raha hai...")
                            continue  # Ye wapas upar ja kar same link ko dobara try karega
                        except Exception as refresh_err:
                            await m.reply_text(f"🚫 **Auto-Fix Failed:**\nKya jis account ka SESSION STRING dala hai, wo sach me is group me join hai? Error: {refresh_err}")
                            break
                    else:
                        await m.reply_text(f"🚫 **Channel Access Error:** `{e}`")
                    break

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

                chunk_has_media = False

                for msg in messages:
                    if STOP_REQUESTED: break
                    if not msg or msg.empty: continue
                        
                    if msg.text and any(k in msg.text.lower() for k in stop_keywords):
                        STOP_REQUESTED = True
                        await m.reply_text(f"✅ **'Subject/Chapter Complete' detected!** Loop automatically stopped.")
                        break

                    target_msg = None
                    if getattr(msg, "media", None) and not getattr(msg, "web_page", None):
                        target_msg = msg
                    elif getattr(msg, "reply_to_message", None) and getattr(msg.reply_to_message, "media", None):
                        target_msg = msg.reply_to_message

                    if not target_msg: continue

                    if getattr(target_msg, "document", None) or getattr(target_msg, "video", None) or getattr(target_msg, "audio", None) or getattr(target_msg, "photo", None):
                        chunk_has_media = True
                        no_media_count = 0
                        
                        is_video = False
                        is_pdf = False
                        ext = "pdf"
                        
                        if target_msg.video or getattr(target_msg, "audio", None):
                            is_video = True
                            ext = "mp4" if target_msg.video else "mp3"
                        elif target_msg.document:
                            fname = getattr(target_msg.document, "file_name", "") or ""
                            mime = getattr(target_msg.document, "mime_type", "") or ""
                            try: ext = fname.split(".")[-1].lower()
                            except: ext = "pdf"
                            
                            if ext in ['mp4', 'mkv', 'webm', 'ts'] or "video" in mime.lower():
                                is_video = True
                            else:
                                is_pdf = True
                        else:
                            is_pdf = True 
                            
                        file_type_text = "Video 🎬" if is_video else "PDF 📄"
                        
                        caption = msg.text if msg.text else (target_msg.caption or "")
                        batch_name, sub_name, chap_name, item_title = parse_caption(caption, is_pdf=is_pdf)

                        header_card = (
                            f"🔄 **CURRENTLY PROCESSING (MSG #{target_msg.id})**\n"
                            f"━━━━━━━━━━━━━━━━━━\n"
                            f"🏷 **Batch:** `{batch_name}`\n"
                            f"📁 **Subject:** `{sub_name}`\n"
                            f"📂 **Chapter:** `{chap_name}`\n"
                            f"📌 **Type:** {file_type_text}\n"
                            f"▶️ **Title:** `{item_title}`"
                        )

                        unique_id = uuid.uuid4().hex[:5]
                        safe_file_name = f"AdminData_{target_msg.id}_{unique_id}.{ext}"
                        save_path = f"./{safe_file_name}"

                        try:
                            # 🔥 AUTO RETRY SYSTEM FOR DOWNLOAD (0 Bytes Fix) 🔥
                            max_retries = 3
                            download_success = False
                            
                            for attempt in range(max_retries):
                                try:
                                    start_t = time.time()
                                    await status.edit(f"{header_card}\n━━━━━━━━━━━━━━━━━━\n⬇️ **Status:** Downloading... (Attempt {attempt+1}/3)")
                                    
                                    await userbot.download_media(target_msg, file_name=save_path, progress=progress, progress_args=(status, start_t, header_card))

                                    if not os.path.exists(save_path) or os.path.getsize(save_path) == 0:
                                        raise Exception("0 bytes corrupted file")
                                    
                                    download_success = True
                                    break
                                except FloodWait as e:
                                    await asyncio.sleep(e.value)
                                except Exception as e:
                                    if os.path.exists(save_path): os.remove(save_path)
                                    if attempt < max_retries - 1:
                                        await asyncio.sleep(2) # 2 second wait before retry
                                    else:
                                        raise Exception(f"Failed to download after {max_retries} attempts.")

                            if not download_success:
                                raise Exception("Download fully failed.")

                            # 🔥 UPLOAD PROCESS 🔥
                            final_link = ""
                            try:
                                await status.edit(f"{header_card}\n━━━━━━━━━━━━━━━━━━\n☁️ **Status:** Uploading to Archive.org...")
                                upload_kwargs = {"access_key": IA_ACCESS_KEY, "secret_key": IA_SECRET_KEY}
                                await asyncio.to_thread(upload, IA_BUCKET, files={safe_file_name: save_path}, **upload_kwargs)
                                final_link = f"https://archive.org/download/{IA_BUCKET}/{safe_file_name}"
                            
                            except Exception as upload_err:
                                if is_pdf and "unacceptable" in str(upload_err).lower():
                                    await status.edit(f"{header_card}\n━━━━━━━━━━━━━━━━━━\n⚠️ **Archive.org Rejected PDF!**\n🔄 Uploading to Backup Server...")
                                    
                                    def backup_upload():
                                        with open(save_path, 'rb') as f:
                                            res = requests.post("https://jitendra55566-bot-uplod.hf.space/api/upload", files={'file': f})
                                        if res.status_code == 200 and 'file_url' in res.json():
                                            return res.json()['file_url']
                                        else:
                                            raise Exception(f"Backup Server Failed: {res.text}")
                                    
                                    final_link = await asyncio.to_thread(backup_upload)
                                else:
                                    raise Exception(f"Upload Failed: {upload_err}")

                            await status.edit(f"{header_card}\n━━━━━━━━━━━━━━━━━━\n🔄 **Status:** Updating Firebase Database...")

                            if not FIREBASE_URL:
                                raise Exception("FIREBASE_URL Variable Empty hai. HuggingFace Secrets check karein.")

                            cat_id = get_id(batch_name)
                            sub_id = get_id(f"{batch_name}_{sub_name}")
                            
                            chap_url = f"{FIREBASE_URL}/categories/{cat_id}/subcategories/{sub_id}/subcategories"
                            chap_id = await asyncio.to_thread(resolve_fb_id, chap_url, chap_name, target_msg.id)
                            
                            title_hash = get_id(f"{chap_name}_{item_title}")[:6]
                            item_id = f"z_{target_msg.id:08d}_{title_hash}"
                            
                            base_match_title = get_base_title(item_title)
                            memory_key = f"{chap_id}_{base_match_title}"

                            def push_to_firebase():
                                patch_firebase(f"{FIREBASE_URL}/categories/{cat_id}.json", {"title": batch_name, "show_on_home": True, "show_on_batches": True})
                                patch_firebase(f"{FIREBASE_URL}/categories/{cat_id}/subcategories/{sub_id}.json", {"title": sub_name})
                                patch_firebase(f"{FIREBASE_URL}/categories/{cat_id}/subcategories/{sub_id}/subcategories/{chap_id}.json", {"title": chap_name})

                                if is_video:
                                    patch_firebase(f"{FIREBASE_URL}/categories/{cat_id}/subcategories/{sub_id}/subcategories/{chap_id}/videos/{item_id}.json", {
                                        "title": item_title, 
                                        "url": final_link, 
                                        "priority": 0, 
                                        "timestamp": target_msg.id
                                    })
                                    video_memory_cache[memory_key] = item_id
                                    
                                else:
                                    pdf_data = {
                                        "title": item_title, 
                                        "url": final_link, 
                                        "allow_download": True, 
                                        "allow_google": True, 
                                        "priority": 0, 
                                        "timestamp": target_msg.id
                                    }

                                    matched_video_id = None
                                    if memory_key in video_memory_cache:
                                        matched_video_id = video_memory_cache[memory_key]
                                    elif last_fallback_video and last_fallback_video["chap_id"] == chap_id:
                                        matched_video_id = last_fallback_video["video_id"]
                                    
                                    if matched_video_id:
                                        patch_firebase(f"{FIREBASE_URL}/categories/{cat_id}/subcategories/{sub_id}/subcategories/{chap_id}/videos/{matched_video_id}/pdfs/{item_id}.json", pdf_data)
                                    else:
                                        patch_firebase(f"{FIREBASE_URL}/categories/{cat_id}/subcategories/{sub_id}/subcategories/{chap_id}/pdfs/{item_id}.json", pdf_data)

                            await asyncio.to_thread(push_to_firebase)

                            if is_video: 
                                last_fallback_video = {"chap_id": chap_id, "video_id": item_id}
                            
                            await m.reply_text(f"✅ **{file_type_text} UPLOADED!**\n🎬 **Title:** {item_title}\n🔢 **MSG ID:** #{target_msg.id}")
                            extracted_count += 1

                        except FloodWait as e:
                            await asyncio.sleep(e.value)
                        except Exception as e:
                            await m.reply_text(f"❌ **Error on MSG #{target_msg.id}:**\n`{e}`")

                        finally:
                            if os.path.exists(save_path): os.remove(save_path)

                if not chunk_has_media:
                    no_media_count += 5
                    try:
                        await status.edit(f"⏩ **Skipping MSG {current_msg_id} to {current_msg_id+4}**\n*(No File Found)*")
                    except FloodWait as e:
                        await asyncio.sleep(e.value)
                    except Exception:
                        pass

                current_msg_id += 5
                
                if no_media_count > 150:
                    await status.edit(f"🛑 **Loop Finished.** 150+ khali messages mile.\nUploaded: `{extracted_count}`")
                    break

                await asyncio.sleep(1)

            await status.delete()

# --- STARTUP ---
async def main():
    threading.Thread(target=run_web_server, daemon=True).start()
    
    global USERBOT_ALIVE
    print("🚀 Starting Main Bot...")
    await bot.start()
    
    if userbot:
        try:
            print("⏳ Starting Userbot...")
            await userbot.start()
            USERBOT_ALIVE = True
            print("🔄 Loading chats to fix Peer ID issue...")
            # 🔥 Bot on hote hi chats memory me save kar lega
            async for dialog in userbot.get_dialogs(limit=50):
                pass
            print("✅ Userbot Online & Chats Loaded!")
        except Exception as e:
            print(f"❌ Userbot Error: {e}")
            USERBOT_ALIVE = False

    print("✅ All Systems Ready!")
    await idle()
    await bot.stop()

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