Blame for bot.py

import os
import asyncio
import logging
import time
import threading
import shutil
from typing import Optional, List, Dict, Any
import aiosqlite
import gradio as gr
from hydrogram import Client, filters
from hydrogram.errors import FloodWait, RPCError
from hydrogram.types import (
    InputMediaPhoto, InputMediaVideo, InputMediaDocument, 
    InputMediaAudio, Message
)
from fastapi import FastAPI
import uvicorn

try:
    from hydrogram.errors import FloodPremiumWait
except ImportError:
    class FloodPremiumWait(FloodWait):
        pass

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s | %(levelname)-7s | %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S'
)
logger = logging.getLogger("UniversalArchiver")

API_ID = int(os.environ.get("API_ID", 0))
API_HASH = os.environ.get("API_HASH", "")
SESSION_STRING = os.environ.get("SESSION_STRING", "")
TARGET_CHAT_RAW = os.environ.get("TARGET_CHAT", "me")
AUTO_SCAN_ALL = os.environ.get("AUTO_SCAN_ALL", "1") == "1"

DATA_DIR = "/tmp/data"
DOWNLOAD_DIR = os.path.join(DATA_DIR, "downloads")
DB_PATH = os.path.join(DATA_DIR, "archiver_v7.db")
os.makedirs(DOWNLOAD_DIR, exist_ok=True)

app = Client(
    "universal_archiver",
    api_id=API_ID,
    api_hash=API_HASH,
    session_string=SESSION_STRING,
    in_memory=True,
    workers=8,
    max_concurrent_transmissions=4,
    device_model="Archiver Pro",
    system_version="Linux 64bit",
    app_version="7.0.0",
    lang_code="ar"
)

TARGET_CHAT = "me"
TARGET_CHAT_ID: Optional[int] = None
MAIN_LOOP: Optional[asyncio.AbstractEventLoop] = None
KNOWN_CHATS: Dict[int, Dict[str, Any]] = {}

stats = {
    "total_queued": 0,
    "total_transferred": 0,
    "total_failed": 0,
    "discovered_chats": 0,
    "active_scans": 0,
    "status": "Starting..."
}

class SmartFloodManager:
    def __init__(self, normal_delay: float = 1.5):
        self.normal_delay = normal_delay
        self.last_call = 0.0
        self.lock = asyncio.Lock()
        self.flood_until = 0.0

    async def wait(self):
        async with self.lock:
            now = time.time()
            if self.flood_until > now:
                sleep_time = self.flood_until - now
                logger.warning(f"FloodWait active, sleeping {sleep_time:.1f}s...")
                await asyncio.sleep(sleep_time)
            
            now = time.time()
            elapsed = now - self.last_call
            if elapsed < self.normal_delay:
                await asyncio.sleep(self.normal_delay - elapsed)
            self.last_call = time.time()

    def handle_flood(self, seconds: int):
        wait_seconds = max(seconds, 5) + 2
        self.flood_until = time.time() + wait_seconds
        logger.warning(f"FloodWait received: {seconds}s. Cooling down for {wait_seconds}s.")

flood_manager = SmartFloodManager(normal_delay=1.5)

class SafeDBManager:
    def __init__(self, db_path: str):
        self.db_path = db_path
        self.lock = asyncio.Lock()
        self.conn: Optional[aiosqlite.Connection] = None

    async def connect(self):
        async with self.lock:
            if self.conn:
                try:
                    await self.conn.close()
                except Exception:
                    pass
                self.conn = None

            if os.path.exists(self.db_path):
                try:
                    test_conn = await aiosqlite.connect(self.db_path)
                    cur = await test_conn.execute("PRAGMA integrity_check;")
                    res = await cur.fetchone()
                    await test_conn.close()
                    if not res or res[0] != "ok":
                        raise Exception("Corrupt DB check failed")
                except Exception as e:
                    logger.error(f"Resetting corrupt DB: {e}")
                    for ext in ["", "-wal", "-shm", "-journal"]:
                        f = f"{self.db_path}{ext}"
                        if os.path.exists(f):
                            try:
                                os.remove(f)
                            except Exception:
                                pass

            self.conn = await aiosqlite.connect(self.db_path)
            self.conn.row_factory = aiosqlite.Row
            await self.conn.execute("PRAGMA journal_mode=WAL;")
            await self.conn.execute("PRAGMA synchronous=NORMAL;")
            await self.conn.execute("PRAGMA busy_timeout=60000;")

            await self.conn.execute("""
                CREATE TABLE IF NOT EXISTS jobs (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    chat_id INTEGER NOT NULL,
                    message_id INTEGER NOT NULL,
                    media_group_id TEXT,
                    media_type TEXT,
                    caption TEXT,
                    source_name TEXT,
                    source_username TEXT,
                    album_total INTEGER DEFAULT 1,
                    album_index INTEGER DEFAULT 1,
                    status TEXT DEFAULT 'pending',
                    attempts INTEGER DEFAULT 0,
                    error TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    UNIQUE(chat_id, message_id)
                )
            """)
            await self.conn.execute("CREATE INDEX IF NOT EXISTS idx_jobs_status ON jobs(status)")
            await self.conn.execute("CREATE INDEX IF NOT EXISTS idx_jobs_media_group ON jobs(media_group_id)")

            await self.conn.execute("""
                CREATE TABLE IF NOT EXISTS scrape_checkpoints (
                    chat_id INTEGER PRIMARY KEY,
                    chat_title TEXT,
                    chat_type TEXT,
                    last_scraped_id INTEGER DEFAULT 0,
                    total_queued INTEGER DEFAULT 0,
                    status TEXT DEFAULT 'idle',
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)
            await self.conn.commit()
            logger.info("Database initialized with single connection lock.")

    async def execute(self, query: str, params: tuple = ()) -> aiosqlite.Cursor:
        async with self.lock:
            try:
                cur = await self.conn.execute(query, params)
                await self.conn.commit()
                return cur
            except aiosqlite.DatabaseError as e:
                logger.error(f"Database error during query, auto-recovering: {e}")
                await self.connect()
                cur = await self.conn.execute(query, params)
                await self.conn.commit()
                return cur

    async def fetchone(self, query: str, params: tuple = ()):
        async with self.lock:
            try:
                cur = await self.conn.execute(query, params)
                row = await cur.fetchone()
                return dict(row) if row else None
            except aiosqlite.DatabaseError as e:
                logger.error(f"Database error during fetchone, auto-recovering: {e}")
                await self.connect()
                cur = await self.conn.execute(query, params)
                row = await cur.fetchone()
                return dict(row) if row else None

    async def fetchall(self, query: str, params: tuple = ()):
        async with self.lock:
            try:
                cur = await self.conn.execute(query, params)
                rows = await cur.fetchall()
                return [dict(r) for r in rows]
            except aiosqlite.DatabaseError as e:
                logger.error(f"Database error during fetchall, auto-recovering: {e}")
                await self.connect()
                cur = await self.conn.execute(query, params)
                rows = await cur.fetchall()
                return [dict(r) for r in rows]

db = SafeDBManager(DB_PATH)

async def add_job_to_queue(message: Message, album_info: Optional[Dict[str, int]] = None) -> bool:
    if not message.media:
        return False
    
    chat_id = int(message.chat.id)
    message_id = int(message.id)
    media_group_id = str(message.media_group_id) if message.media_group_id else None
    
    media_type = "unknown"
    if message.photo: media_type = "photo"
    elif message.video: media_type = "video"
    elif message.document: media_type = "document"
    elif message.audio: media_type = "audio"
    elif message.voice: media_type = "voice"
    elif message.video_note: media_type = "video_note"
    elif message.animation: media_type = "animation"

    chat = message.chat
    source_name = chat.title or chat.first_name or chat.username or f"Chat_{chat_id}"
    source_username = chat.username or ""
    raw_caption = message.caption or ""

    album_total = 1
    album_index = 1
    if album_info:
        album_total = album_info.get("total", 1)
        album_index = album_info.get("index", 1)

    try:
        cur = await db.execute("""
            INSERT OR IGNORE INTO jobs 
            (chat_id, message_id, media_group_id, media_type, caption, source_name, source_username, album_total, album_index, status)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, 'pending')
        """, (chat_id, message_id, media_group_id, media_type, raw_caption, source_name, source_username, album_total, album_index))
        if cur.rowcount > 0:
            stats["total_queued"] += 1
            return True
        return False
    except Exception as e:
        logger.error(f"Error queuing message {message_id}: {e}")
        return False

def format_smart_caption(raw_cap: str, source_name: str, source_username: str, msg_id: int, album_index: int, album_total: int) -> str:
    lines = []
    if raw_cap and raw_cap.strip():
        lines.append(raw_cap.strip())
        lines.append("")
    
    if source_username:
        source_tag = f"[{source_name}](https://t.me/{source_username}/{msg_id})"
    else:
        source_tag = f"**{source_name}**"

    lines.append(f"📢 المصدر: {source_tag}")
    
    if album_total > 1:
        lines.append(f"🔢 المقطع: {album_index} من إجمالي {album_total}")

    return "\n".join(lines)

async def process_job(job: Dict[str, Any]) -> str:
    chat_id = job["chat_id"]
    msg_id = job["message_id"]
    raw_caption = job["caption"]
    source_name = job["source_name"]
    source_username = job.get("source_username", "")
    album_index = job["album_index"]
    album_total = job["album_total"]
    mtype = job["media_type"]
    
    caption = format_smart_caption(raw_caption, source_name, source_username, msg_id, album_index, album_total)
    await flood_manager.wait()

    try:
        await app.copy_message(
            chat_id=TARGET_CHAT,
            from_chat_id=chat_id,
            message_id=msg_id,
            caption=caption
        )
        logger.info(f"⚡ Fast copied: {msg_id} from [{source_name}]")
        return "SUCCESS"
    except (FloodWait, FloodPremiumWait) as e:
        flood_manager.handle_flood(getattr(e, 'value', 30))
        raise e
    except Exception:
        pass

    downloaded_file = None
    try:
        msg = await app.get_messages(chat_id, msg_id)
        if not msg or not msg.media:
            return "DELETED"

        ext_map = {
            "photo": ".jpg",
            "video": ".mp4",
            "animation": ".mp4",
            "video_note": ".mp4",
            "audio": ".mp3",
            "voice": ".ogg",
            "document": ""
        }
        ext = ext_map.get(mtype, "")
        target_path = os.path.join(DOWNLOAD_DIR, f"media_{chat_id}_{msg_id}{ext}")
        
        downloaded = await app.download_media(msg, file_name=target_path)
        if not downloaded or not os.path.exists(downloaded):
            raise Exception("Download failed")
        downloaded_file = downloaded

        if mtype == "photo":
            try:
                await app.send_photo(TARGET_CHAT, photo=downloaded_file, caption=caption)
            except Exception:
                await app.send_document(TARGET_CHAT, document=downloaded_file, caption=caption)
        elif mtype == "video":
            await app.send_video(TARGET_CHAT, video=downloaded_file, caption=caption, supports_streaming=True)
        elif mtype == "document":
            await app.send_document(TARGET_CHAT, document=downloaded_file, caption=caption)
        elif mtype == "audio":
            await app.send_audio(TARGET_CHAT, audio=downloaded_file, caption=caption)
        elif mtype == "voice":
            await app.send_voice(TARGET_CHAT, voice=downloaded_file, caption=caption)
        elif mtype == "video_note":
            await app.send_video_note(TARGET_CHAT, video_note=downloaded_file)
        elif mtype == "animation":
            await app.send_animation(TARGET_CHAT, animation=downloaded_file, caption=caption)
        else:
            await app.send_document(TARGET_CHAT, document=downloaded_file, caption=caption)

        logger.info(f"📥 Downloaded & uploaded: {msg_id} from [{source_name}]")
        return "SUCCESS"
    except (FloodWait, FloodPremiumWait) as e:
        flood_manager.handle_flood(getattr(e, 'value', 30))
        raise e
    except Exception as e:
        logger.error(f"❌ Failed processing {msg_id}: {e}")
        return "RETRY"
    finally:
        if downloaded_file and os.path.exists(downloaded_file):
            try:
                if os.path.isdir(downloaded_file):
                    shutil.rmtree(downloaded_file, ignore_errors=True)
                else:
                    os.remove(downloaded_file)
            except Exception:
                pass

async def queue_worker_loop():
    logger.info("⚙️ Queue worker loop started")
    while True:
        try:
            if not app.is_connected:
                await asyncio.sleep(5)
                continue

            job = await db.fetchone("""
                SELECT * FROM jobs 
                WHERE status = 'pending' AND attempts < 10
                ORDER BY attempts ASC, id ASC LIMIT 1
            """)
            if not job:
                await asyncio.sleep(2)
                continue

            await db.execute("UPDATE jobs SET status='processing' WHERE id=?", (job["id"],))
            result = await process_job(job)

            if result == "SUCCESS":
                await db.execute("UPDATE jobs SET status='completed', error=NULL WHERE id=?", (job["id"],))
                stats["total_transferred"] += 1
            elif result == "DELETED":
                await db.execute("UPDATE jobs SET status='failed', error='Deleted' WHERE id=?", (job["id"],))
                stats["total_failed"] += 1
            else:
                await db.execute("""
                    UPDATE jobs SET status='pending', attempts=attempts+1, error='Retry' WHERE id=?
                """, (job["id"],))

        except (FloodWait, FloodPremiumWait) as e:
            sec = getattr(e, 'value', 30)
            logger.warning(f"FloodWait sleeping {sec}s")
            await asyncio.sleep(sec + 1)
        except Exception as e:
            logger.error(f"Worker exception: {e}")
            await asyncio.sleep(4)

async def warm_up_dialogs_cache() -> List[Dict[str, Any]]:
    logger.info("🔄 Warming up dialogs cache...")
    discovered = []
    try:
        async for dialog in app.get_dialogs():
            chat = dialog.chat
            if not chat:
                continue

            if TARGET_CHAT_ID and chat.id == TARGET_CHAT_ID:
                continue
            if str(chat.id) == str(TARGET_CHAT):
                continue

            type_str = str(chat.type).lower()
            if any(k in type_str for k in ["channel", "supergroup", "group"]):
                title = chat.title or chat.first_name or f"Chat_{chat.id}"
                chat_info = {
                    "id": chat.id,
                    "title": title,
                    "username": chat.username or "",
                    "type": type_str
                }
                KNOWN_CHATS[chat.id] = chat_info
                discovered.append(chat_info)
                logger.info(f"📌 Found: [{title}] ({chat.id}) - {type_str}")

        stats["discovered_chats"] = len(discovered)
        logger.info(f"✅ Discovered {len(discovered)} chats/channels")
        return discovered
    except Exception as e:
        logger.error(f"Dialogs error: {e}")
        return []

async def scrape_single_chat(chat_id: int, chat_title: str, limit: int = 0) -> int:
    try:
        last_row = await db.fetchone("SELECT last_scraped_id FROM scrape_checkpoints WHERE chat_id=?", (chat_id,))
        last_id = last_row["last_scraped_id"] if last_row else 0

        logger.info(f"🔍 Scraping: {chat_title} (from msg {last_id})...")
        count = 0
        media_groups: Dict[str, List[Message]] = {}
        highest_msg_id = last_id

        async for message in app.get_chat_history(chat_id, limit=limit if limit > 0 else 0):
            if message.id <= last_id and last_id > 0:
                break
            
            if message.id > highest_msg_id:
                highest_msg_id = message.id

            if message.media:
                if message.media_group_id:
                    mg_id = str(message.media_group_id)
                    media_groups.setdefault(mg_id, []).append(message)
                else:
                    if await add_job_to_queue(message):
                        count += 1

            if count > 0 and count % 100 == 0:
                await db.execute("""
                    INSERT INTO scrape_checkpoints (chat_id, chat_title, last_scraped_id, total_queued, status)
                    VALUES (?, ?, ?, ?, 'in_progress')
                    ON CONFLICT(chat_id) DO UPDATE SET 
                        last_scraped_id=excluded.last_scraped_id,
                        total_queued=total_queued+excluded.total_queued,
                        updated_at=CURRENT_TIMESTAMP
                """, (chat_id, chat_title, highest_msg_id, 100))
                await asyncio.sleep(0.1)

        for mg_id, msgs in media_groups.items():
            msgs.sort(key=lambda m: m.id)
            total = len(msgs)
            for idx, msg in enumerate(msgs, start=1):
                if await add_job_to_queue(msg, album_info={"total": total, "index": idx}):
                    count += 1

        await db.execute("""
            INSERT INTO scrape_checkpoints (chat_id, chat_title, last_scraped_id, total_queued, status)
            VALUES (?, ?, ?, ?, 'completed')
            ON CONFLICT(chat_id) DO UPDATE SET 
                last_scraped_id=excluded.last_scraped_id,
                total_queued=total_queued+excluded.total_queued,
                status='completed',
                updated_at=CURRENT_TIMESTAMP
        """, (chat_id, chat_title, highest_msg_id, count))

        logger.info(f"✅ Finished {chat_title}: queued {count} files.")
        return count
    except Exception as e:
        logger.warning(f"Skipping {chat_title} ({chat_id}): {e}")
        return 0

async def scrape_all_joined_chats(limit_per_chat: int = 0) -> str:
    stats["active_scans"] += 1
    try:
        chats = await warm_up_dialogs_cache()
        if not chats:
            return "No chats found."

        logger.info(f"🚀 Scraping all {len(chats)} chats...")
        total_all = 0
        for idx, chat in enumerate(chats, 1):
            c_id = chat["id"]
            c_title = chat["title"]
            logger.info(f"[{idx}/{len(chats)}] Scraping: {c_title}...")
            queued = await scrape_single_chat(c_id, c_title, limit=limit_per_chat)
            total_all += queued
            await asyncio.sleep(1.0)

        msg = f"Completed all chats ({len(chats)}). Queued {total_all} files."
        logger.info(msg)
        return msg
    except Exception as e:
        err = f"Scrape all error: {e}"
        logger.error(err)
        return err
    finally:
        stats["active_scans"] = max(0, stats["active_scans"] - 1)

@app.on_message(filters.incoming & ~filters.me & filters.media)
async def realtime_all_chats_handler(client, message: Message):
    try:
        if TARGET_CHAT_ID and message.chat.id == TARGET_CHAT_ID:
            return
        if str(message.chat.id) == str(TARGET_CHAT):
            return

        type_str = str(message.chat.type).lower()
        if message.chat.id not in KNOWN_CHATS and any(k in type_str for k in ["channel", "supergroup", "group"]):
            KNOWN_CHATS[message.chat.id] = {
                "id": message.chat.id,
                "title": message.chat.title or f"Chat_{message.chat.id}",
                "username": message.chat.username or "",
                "type": type_str
            }

        if message.media_group_id:
            await asyncio.sleep(2.0)
            try:
                album_msgs = await client.get_media_group(message.chat.id, message.id)
                album_msgs.sort(key=lambda m: m.id)
                total = len(album_msgs)
                for idx, msg in enumerate(album_msgs, start=1):
                    await add_job_to_queue(msg, album_info={"total": total, "index": idx})
            except Exception:
                await add_job_to_queue(message)
        else:
            await add_job_to_queue(message)
            
        logger.info(f"✨ Realtime media: [{message.chat.title or message.chat.id}]")
    except Exception as e:
        logger.error(f"Realtime error: {e}")

async def resolve_target_chat():
    global TARGET_CHAT, TARGET_CHAT_ID
    raw = TARGET_CHAT_RAW.strip()
    if raw.lower() in ["me", "saved", ""]:
        TARGET_CHAT = "me"
        TARGET_CHAT_ID = None
        logger.info("Target: Saved Messages (me)")
        return
    try:
        chat_id_int = int(raw)
        chat = await app.get_chat(chat_id_int)
        TARGET_CHAT = chat_id_int
        TARGET_CHAT_ID = chat.id
        logger.info(f"Target: {chat.title} ({TARGET_CHAT})")
    except ValueError:
        try:
            chat = await app.get_chat(raw)
            TARGET_CHAT = chat.id
            TARGET_CHAT_ID = chat.id
            logger.info(f"Target: {chat.title} ({TARGET_CHAT})")
        except Exception as e:
            logger.error(f"Target error ({raw}): {e}")
            TARGET_CHAT = "me"

def create_dashboard():
    with gr.Blocks(title="Telegram Archiver") as demo:
        gr.Markdown("# 🌐 Universal Telegram Archiver Pro")

        with gr.Row():
            gr.Textbox(label="Status", value="Active ✅", interactive=False)
            gr.Number(label="Discovered Chats", value=0, interactive=False)
            gr.Number(label="Queue Size", value=0, interactive=False)
            gr.Number(label="Transferred", value=0, interactive=False)

        with gr.Row():
            scan_all_btn = gr.Button("🔄 Scan & Scrape All Joined Chats", variant="primary", scale=2)
            limit_input = gr.Number(label="Limit per chat (0 = All)", value=0, scale=1)

        scrape_status = gr.Textbox(label="Operation Result", interactive=False)

        def trigger_scan_all(lim):
            if MAIN_LOOP and MAIN_LOOP.is_running():
                asyncio.run_coroutine_threadsafe(
                    scrape_all_joined_chats(int(lim)), 
                    MAIN_LOOP
                )
                return "Started full scan in background!"
            return "Bot is not connected."

        scan_all_btn.click(fn=trigger_scan_all, inputs=[limit_input], outputs=[scrape_status])

    return demo

def run_web_server():
    fastapi_app = FastAPI()
    
    @fastapi_app.get("/health")
    async def health():
        return {"status": "ok", "stats": stats}

    @fastapi_app.get("/ping")
    async def ping():
        return {"status": "alive", "stats": stats}

    demo = create_dashboard()
    gr.mount_gradio_app(fastapi_app, demo, path="/")
    
    uvicorn.run(fastapi_app, host="0.0.0.0", port=7860)

async def main():
    global MAIN_LOOP
    MAIN_LOOP = asyncio.get_running_loop()

    logger.info("Connecting to Telegram...")
    await app.start()
    
    me = await app.get_me()
    logger.info(f"Logged in: {me.first_name} (@{me.username or me.id})")

    await resolve_target_chat()
    await db.connect()

    asyncio.create_task(queue_worker_loop())

    discovered = await warm_up_dialogs_cache()
    logger.info(f"Ready: {len(discovered)} chats.")

    if AUTO_SCAN_ALL:
        logger.info("Auto-scanning all chats in background...")
        asyncio.create_task(scrape_all_joined_chats(limit_per_chat=0))

    stats["status"] = "Active"
    
    stop_event = asyncio.Event()
    await stop_event.wait()

if __name__ == "__main__":
    web_thread = threading.Thread(target=run_web_server, daemon=True)
    web_thread.start()

    loop = asyncio.get_event_loop()
    try:
        loop.run_until_complete(main())
    except KeyboardInterrupt:
        pass
    finally:
        if app.is_connected:
            loop.run_until_complete(app.stop())