ziper/bot2.py

import asyncio
import logging
import os
import shutil
import sqlite3
import subprocess
import time
import uuid
from pathlib import Path
from typing import Optional, List, Tuple

from dotenv import load_dotenv
from PIL import Image, ImageOps, UnidentifiedImageError
from telegram import Update, InputMediaPhoto
from telegram.error import TimedOut, NetworkError, RetryAfter, BadRequest
from telegram.ext import (
    Application,
    ApplicationBuilder,
    CommandHandler,
    ContextTypes,
    MessageHandler,
    filters,
)

load_dotenv()

BOT_TOKEN = os.getenv("BOT_TOKEN", "").strip()
FFMPEG_PATH = os.getenv("FFMPEG_PATH", "/usr/bin/ffmpeg").strip()

MAX_MEDIA_PER_FINAL = int(os.getenv("MAX_MEDIA_PER_FINAL", "1000"))
STATUS_EVERY = int(os.getenv("STATUS_EVERY", "25"))
SEND_DELAY = float(os.getenv("SEND_DELAY", "1.0"))
DOWNLOAD_DELAY = float(os.getenv("DOWNLOAD_DELAY", "0.05"))

BASE_DIR = Path(__file__).resolve().parent
DATA_DIR = BASE_DIR / "data"
DOWNLOAD_DIR = BASE_DIR / "downloads"
OUTPUT_DIR = BASE_DIR / "output"
DB_PATH = DATA_DIR / "queue.sqlite3"

for directory in [DATA_DIR, DOWNLOAD_DIR, OUTPUT_DIR]:
    directory.mkdir(parents=True, exist_ok=True)

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
)
logger = logging.getLogger("media_to_photo_bot")


def db() -> sqlite3.Connection:
    conn = sqlite3.connect(DB_PATH)
    conn.execute("PRAGMA journal_mode=WAL")
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS inbox_media (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            chat_id INTEGER NOT NULL,
            user_id INTEGER NOT NULL,
            message_id INTEGER NOT NULL,
            file_id TEXT NOT NULL,
            file_unique_id TEXT,
            media_type TEXT NOT NULL,
            file_name TEXT,
            mime_type TEXT,
            queued INTEGER NOT NULL DEFAULT 0,
            processed INTEGER NOT NULL DEFAULT 0,
            created_at INTEGER NOT NULL,
            UNIQUE(chat_id, user_id, message_id, file_unique_id)
        )
        """
    )
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS media_queue (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            inbox_id INTEGER,
            chat_id INTEGER NOT NULL,
            user_id INTEGER NOT NULL,
            message_id INTEGER,
            file_id TEXT NOT NULL,
            file_unique_id TEXT,
            media_type TEXT NOT NULL,
            file_name TEXT,
            mime_type TEXT,
            created_at INTEGER NOT NULL,
            processed INTEGER NOT NULL DEFAULT 0,
            UNIQUE(chat_id, user_id, file_unique_id, message_id)
        )
        """
    )

    # Migration for old DBs
    cols = [r[1] for r in conn.execute("PRAGMA table_info(media_queue)").fetchall()]
    if "processed" not in cols:
        conn.execute("ALTER TABLE media_queue ADD COLUMN processed INTEGER NOT NULL DEFAULT 0")
    if "inbox_id" not in cols:
        conn.execute("ALTER TABLE media_queue ADD COLUMN inbox_id INTEGER")
    if "message_id" not in cols:
        conn.execute("ALTER TABLE media_queue ADD COLUMN message_id INTEGER")
    if "file_unique_id" not in cols:
        conn.execute("ALTER TABLE media_queue ADD COLUMN file_unique_id TEXT")

    conn.commit()
    return conn


def extract_media(update: Update):
    msg = update.effective_message
    if not msg:
        return None

    file_id = None
    file_unique_id = None
    media_type = None
    file_name = None
    mime_type = None

    if msg.photo:
        best_photo = msg.photo[-1]
        file_id = best_photo.file_id
        file_unique_id = best_photo.file_unique_id
        media_type = "image"
        file_name = f"telegram_photo_{file_unique_id}.jpg"
        mime_type = "image/jpeg"

    elif msg.video:
        file_id = msg.video.file_id
        file_unique_id = msg.video.file_unique_id
        media_type = "video"
        file_name = msg.video.file_name or f"telegram_video_{file_unique_id}.mp4"
        mime_type = msg.video.mime_type or "video/mp4"

    elif msg.animation:
        file_id = msg.animation.file_id
        file_unique_id = msg.animation.file_unique_id
        media_type = "video"
        file_name = msg.animation.file_name or f"telegram_animation_{file_unique_id}.mp4"
        mime_type = msg.animation.mime_type or "video/mp4"

    elif msg.document:
        doc = msg.document
        file_id = doc.file_id
        file_unique_id = doc.file_unique_id
        file_name = doc.file_name or f"telegram_document_{file_unique_id}"
        mime_type = doc.mime_type or ""

        if mime_type.startswith("image/"):
            media_type = "image"
        elif mime_type.startswith("video/"):
            media_type = "video"
        else:
            return None

    if not file_id or not media_type:
        return None

    return {
        "chat_id": update.effective_chat.id,
        "user_id": update.effective_user.id,
        "message_id": msg.message_id,
        "file_id": file_id,
        "file_unique_id": file_unique_id or file_id,
        "media_type": media_type,
        "file_name": file_name,
        "mime_type": mime_type,
    }


def save_to_inbox(media: dict) -> Optional[int]:
    conn = db()
    cur = conn.cursor()
    try:
        cur.execute(
            """
            INSERT OR IGNORE INTO inbox_media
            (chat_id, user_id, message_id, file_id, file_unique_id, media_type, file_name, mime_type, created_at)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                media["chat_id"],
                media["user_id"],
                media["message_id"],
                media["file_id"],
                media["file_unique_id"],
                media["media_type"],
                media["file_name"],
                media["mime_type"],
                int(time.time()),
            ),
        )
        conn.commit()

        row = conn.execute(
            """
            SELECT id FROM inbox_media
            WHERE chat_id = ? AND user_id = ? AND message_id = ? AND file_unique_id = ?
            """,
            (
                media["chat_id"],
                media["user_id"],
                media["message_id"],
                media["file_unique_id"],
            ),
        ).fetchone()

        return row[0] if row else None
    finally:
        conn.close()


def add_to_queue(media: dict, inbox_id: Optional[int] = None) -> bool:
    conn = db()
    try:
        cur = conn.cursor()
        cur.execute(
            """
            INSERT OR IGNORE INTO media_queue
            (inbox_id, chat_id, user_id, message_id, file_id, file_unique_id, media_type, file_name, mime_type, created_at)
            VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                inbox_id,
                media["chat_id"],
                media["user_id"],
                media.get("message_id"),
                media["file_id"],
                media["file_unique_id"],
                media["media_type"],
                media["file_name"],
                media["mime_type"],
                int(time.time()),
            ),
        )

        inserted = cur.rowcount > 0

        if inbox_id:
            conn.execute("UPDATE inbox_media SET queued = 1 WHERE id = ?", (inbox_id,))

        conn.commit()
        return inserted
    finally:
        conn.close()


def scan_inbox_to_queue(chat_id: int, user_id: int) -> int:
    conn = db()
    try:
        rows = conn.execute(
            """
            SELECT id, chat_id, user_id, message_id, file_id, file_unique_id, media_type, file_name, mime_type
            FROM inbox_media
            WHERE chat_id = ?
              AND user_id = ?
              AND queued = 0
              AND processed = 0
            ORDER BY id ASC
            """
            ,
            (chat_id, user_id),
        ).fetchall()

        added = 0

        for row in rows:
            inbox_id, r_chat_id, r_user_id, message_id, file_id, file_unique_id, media_type, file_name, mime_type = row

            cur = conn.execute(
                """
                INSERT OR IGNORE INTO media_queue
                (inbox_id, chat_id, user_id, message_id, file_id, file_unique_id, media_type, file_name, mime_type, created_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    inbox_id,
                    r_chat_id,
                    r_user_id,
                    message_id,
                    file_id,
                    file_unique_id,
                    media_type,
                    file_name,
                    mime_type,
                    int(time.time()),
                ),
            )

            if cur.rowcount > 0:
                added += 1

            conn.execute("UPDATE inbox_media SET queued = 1 WHERE id = ?", (inbox_id,))

        conn.commit()
        return added
    finally:
        conn.close()


def get_queue(chat_id: int, user_id: int) -> List[Tuple]:
    conn = db()
    try:
        rows = conn.execute(
            """
            SELECT id, inbox_id, file_id, media_type, file_name, mime_type
            FROM media_queue
            WHERE chat_id = ?
              AND user_id = ?
              AND processed = 0
            ORDER BY id ASC
            LIMIT ?
            """,
            (chat_id, user_id, MAX_MEDIA_PER_FINAL),
        ).fetchall()
        return rows
    finally:
        conn.close()


def queue_count(chat_id: int, user_id: int) -> int:
    conn = db()
    try:
        return conn.execute(
            """
            SELECT COUNT(*) FROM media_queue
            WHERE chat_id = ? AND user_id = ? AND processed = 0
            """,
            (chat_id, user_id),
        ).fetchone()[0]
    finally:
        conn.close()


def inbox_count(chat_id: int, user_id: int) -> int:
    conn = db()
    try:
        return conn.execute(
            """
            SELECT COUNT(*) FROM inbox_media
            WHERE chat_id = ? AND user_id = ? AND queued = 0 AND processed = 0
            """,
            (chat_id, user_id),
        ).fetchone()[0]
    finally:
        conn.close()


def clear_queue(chat_id: int, user_id: int) -> int:
    conn = db()
    try:
        cur = conn.cursor()
        cur.execute(
            """
            DELETE FROM media_queue
            WHERE chat_id = ? AND user_id = ? AND processed = 0
            """,
            (chat_id, user_id),
        )
        deleted = cur.rowcount

        conn.execute(
            """
            UPDATE inbox_media
            SET queued = 0
            WHERE chat_id = ? AND user_id = ? AND processed = 0
            """,
            (chat_id, user_id),
        )

        conn.commit()
        return deleted
    finally:
        conn.close()


def mark_processed(queue_ids: List[int]) -> None:
    if not queue_ids:
        return

    conn = db()
    try:
        placeholders = ",".join("?" for _ in queue_ids)

        inbox_ids = [
            r[0]
            for r in conn.execute(
                f"SELECT inbox_id FROM media_queue WHERE id IN ({placeholders}) AND inbox_id IS NOT NULL",
                queue_ids,
            ).fetchall()
        ]

        conn.execute(
            f"UPDATE media_queue SET processed = 1 WHERE id IN ({placeholders})",
            queue_ids,
        )

        if inbox_ids:
            inbox_placeholders = ",".join("?" for _ in inbox_ids)
            conn.execute(
                f"UPDATE inbox_media SET processed = 1, queued = 1 WHERE id IN ({inbox_placeholders})",
                inbox_ids,
            )

        conn.commit()
    finally:
        conn.close()


def delete_processed(chat_id: int, user_id: int) -> None:
    conn = db()
    try:
        conn.execute(
            "DELETE FROM media_queue WHERE chat_id = ? AND user_id = ? AND processed = 1",
            (chat_id, user_id),
        )
        conn.commit()
    finally:
        conn.close()


def safe_suffix(file_name: Optional[str], fallback: str) -> str:
    if not file_name:
        return fallback
    suffix = Path(file_name).suffix.lower()
    if not suffix:
        return fallback
    if len(suffix) > 10:
        return fallback
    return suffix


def ffprobe_path() -> str:
    ffmpeg = Path(FFMPEG_PATH)
    possible = ffmpeg.with_name("ffprobe")
    return str(possible) if possible.exists() else "ffprobe"


def get_video_duration(path: Path) -> Optional[float]:
    try:
        result = subprocess.run(
            [
                ffprobe_path(),
                "-v",
                "error",
                "-show_entries",
                "format=duration",
                "-of",
                "default=noprint_wrappers=1:nokey=1",
                str(path),
            ],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            timeout=30,
        )
        if result.returncode != 0:
            return None
        return float(result.stdout.strip())
    except Exception:
        return None


def convert_image_to_jpg(input_path: Path, output_path: Path) -> bool:
    try:
        with Image.open(input_path) as img:
            img = ImageOps.exif_transpose(img)
            if img.mode in ("RGBA", "LA", "P"):
                img = img.convert("RGBA")
                background = Image.new("RGBA", img.size, (255, 255, 255, 255))
                background.alpha_composite(img)
                img = background.convert("RGB")
            else:
                img = img.convert("RGB")

            img.save(output_path, "JPEG", quality=90, optimize=True)
            return True
    except UnidentifiedImageError:
        return False
    except Exception as e:
        logger.exception("Image conversion failed: %s", e)
        return False


def convert_video_to_jpg(input_path: Path, output_path: Path) -> bool:
    duration = get_video_duration(input_path)
    seek_time = 1

    if duration and duration > 3:
        seek_time = max(1, int(duration / 2))

    try:
        result = subprocess.run(
            [
                FFMPEG_PATH,
                "-y",
                "-ss",
                str(seek_time),
                "-i",
                str(input_path),
                "-frames:v",
                "1",
                "-q:v",
                "3",
                str(output_path),
            ],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            timeout=180,
        )
        return result.returncode == 0 and output_path.exists() and output_path.stat().st_size > 0
    except Exception as e:
        logger.exception("Video conversion failed: %s", e)
        return False


async def safe_reply(msg, text: str):
    try:
        return await msg.reply_text(text, read_timeout=60, write_timeout=60, connect_timeout=30, pool_timeout=30)
    except Exception as e:
        logger.warning("safe_reply failed: %s", e)
        return None


async def safe_edit(message, text: str):
    if not message:
        return
    try:
        return await message.edit_text(text, read_timeout=60, write_timeout=60, connect_timeout=30, pool_timeout=30)
    except BadRequest as e:
        if "Message is not modified" not in str(e):
            logger.warning("safe_edit bad request: %s", e)
    except Exception as e:
        logger.warning("safe_edit failed: %s", e)


async def retry_async(func, *args, retries: int = 5, **kwargs):
    delay = 2
    last_error = None

    for attempt in range(1, retries + 1):
        try:
            return await func(*args, **kwargs)
        except RetryAfter as e:
            await asyncio.sleep(int(e.retry_after) + 2)
            last_error = e
        except (TimedOut, NetworkError) as e:
            last_error = e
            logger.warning("Telegram network error attempt %s/%s: %s", attempt, retries, e)
            await asyncio.sleep(delay)
            delay = min(delay * 2, 30)

    raise last_error


async def start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    text = (
        "بات آماده است.\n\n"
        "حالت جدید:\n"
        "- مدیاها ذخیره می‌شوند.\n"
        "- برای وارد کردن مدیاهای ذخیره‌شده به صف: /scan\n"
        "- برای تبدیل و ارسال: /final\n\n"
        "دستورها:\n"
        "/queue - وضعیت صف و inbox\n"
        "/scan - اضافه کردن مدیاهای دیده‌شده قبلی به صف\n"
        "/final - دانلود، تبدیل به عکس، ارسال خروجی\n"
        "/clear - پاک کردن صف فعلی"
    )
    await safe_reply(update.effective_message, text)


async def queue_status(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    chat_id = update.effective_chat.id
    user_id = update.effective_user.id

    q = queue_count(chat_id, user_id)
    inbox = inbox_count(chat_id, user_id)

    await safe_reply(
        update.effective_message,
        f"صف فعلی: {q}\nمدیای ذخیره‌شده ولی صف‌نشده: {inbox}",
    )


async def scan(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    chat_id = update.effective_chat.id
    user_id = update.effective_user.id

    added = scan_inbox_to_queue(chat_id, user_id)
    q = queue_count(chat_id, user_id)

    await safe_reply(
        update.effective_message,
        f"اسکن انجام شد.\nاضافه‌شده به صف: {added}\nصف فعلی: {q}\nبرای اجرا: /final",
    )


async def clear(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    chat_id = update.effective_chat.id
    user_id = update.effective_user.id

    deleted = clear_queue(chat_id, user_id)
    await safe_reply(update.effective_message, f"صف پاک شد. تعداد حذف‌شده: {deleted}")


async def collect_media(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    media = extract_media(update)
    if not media:
        return

    inbox_id = save_to_inbox(media)

    # Auto queue is still enabled for new media.
    # If you want manual-only queue, comment the next line.
    inserted = add_to_queue(media, inbox_id=inbox_id)

    # Do not reply for every file. This was causing floods/timeouts.
    # Send a small confirmation only every STATUS_EVERY new files.
    if inserted:
        q = queue_count(media["chat_id"], media["user_id"])
        if q == 1 or q % STATUS_EVERY == 0:
            await safe_reply(
                update.effective_message,
                f"ذخیره شد. صف فعلی: {q}\nبرای اجرا: /final",
            )


async def final(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    msg = update.effective_message
    chat_id = update.effective_chat.id
    user_id = update.effective_user.id

    rows = get_queue(chat_id, user_id)

    if not rows:
        await safe_reply(msg, "صف خالی است. اگر مدیاها قبلاً به بات رسیده‌اند، اول /scan را بزن.")
        return

    job_id = uuid.uuid4().hex[:12]
    job_download_dir = DOWNLOAD_DIR / f"{chat_id}_{user_id}_{job_id}"
    job_output_dir = OUTPUT_DIR / f"{chat_id}_{user_id}_{job_id}"
    job_download_dir.mkdir(parents=True, exist_ok=True)
    job_output_dir.mkdir(parents=True, exist_ok=True)

    status = await safe_reply(msg, f"شروع شد. تعداد فایل‌ها: {len(rows)}")

    output_images: List[Path] = []
    processed_queue_ids: List[int] = []
    failed = 0

    try:
        for index, row in enumerate(rows, start=1):
            queue_id, inbox_id, file_id, media_type, file_name, mime_type = row

            if index == 1 or index % STATUS_EVERY == 0 or index == len(rows):
                await safe_edit(status, f"درحال دانلود و تبدیل: {index}/{len(rows)}")

            suffix = safe_suffix(file_name, ".bin")
            input_path = job_download_dir / f"{queue_id}{suffix}"
            output_path = job_output_dir / f"{queue_id}.jpg"

            try:
                tg_file = await retry_async(
                    context.bot.get_file,
                    file_id,
                    retries=5,
                    read_timeout=120,
                    write_timeout=120,
                    connect_timeout=60,
                    pool_timeout=60,
                )

                await retry_async(
                    tg_file.download_to_drive,
                    custom_path=str(input_path),
                    retries=5,
                    read_timeout=300,
                    write_timeout=300,
                    connect_timeout=60,
                    pool_timeout=60,
                )

                ok = False

                if media_type == "image":
                    ok = await asyncio.to_thread(convert_image_to_jpg, input_path, output_path)
                elif media_type == "video":
                    ok = await asyncio.to_thread(convert_video_to_jpg, input_path, output_path)

                if ok:
                    output_images.append(output_path)
                    processed_queue_ids.append(queue_id)
                else:
                    failed += 1

            except Exception as e:
                failed += 1
                logger.exception("Failed item queue_id=%s file_id=%s error=%s", queue_id, file_id, e)

            await asyncio.sleep(DOWNLOAD_DELAY)

        if not output_images:
            await safe_edit(status, f"هیچ خروجی تصویری ساخته نشد. ناموفق: {failed}")
            return

        await safe_edit(status, f"تبدیل تمام شد. درحال ارسال {len(output_images)} عکس...")

        sent = 0

        for i in range(0, len(output_images), 10):
            chunk = output_images[i:i + 10]
            media_group = []
            opened_files = []

            try:
                for image_path in chunk:
                    f = open(image_path, "rb")
                    opened_files.append(f)
                    media_group.append(InputMediaPhoto(media=f))

                await retry_async(
                    context.bot.send_media_group,
                    chat_id=chat_id,
                    media=media_group,
                    reply_to_message_id=msg.message_id,
                    retries=5,
                    read_timeout=180,
                    write_timeout=300,
                    connect_timeout=60,
                    pool_timeout=60,
                )

                sent += len(chunk)

                if sent % 50 == 0 or sent == len(output_images):
                    await safe_edit(status, f"ارسال شده: {sent}/{len(output_images)}")

                await asyncio.sleep(SEND_DELAY)

            finally:
                for f in opened_files:
                    try:
                        f.close()
                    except Exception:
                        pass

        mark_processed(processed_queue_ids)
        delete_processed(chat_id, user_id)

        remain = queue_count(chat_id, user_id)

        done_text = (
            f"تمام شد.\n"
            f"ارسال‌شده: {sent}\n"
            f"پردازش‌شده و حذف‌شده از صف: {len(processed_queue_ids)}\n"
            f"ناموفق: {failed}\n"
            f"باقی‌مانده در صف: {remain}"
        )

        await safe_edit(status, done_text)

    except Exception as e:
        logger.exception("Final command failed: %s", e)
        await safe_edit(status, f"خطا هنگام پردازش: {e}")

    finally:
        shutil.rmtree(job_download_dir, ignore_errors=True)
        shutil.rmtree(job_output_dir, ignore_errors=True)


async def error_handler(update: object, context: ContextTypes.DEFAULT_TYPE) -> None:
    logger.exception("Unhandled error: %s", context.error)


def main() -> None:
    if not BOT_TOKEN:
        raise RuntimeError("BOT_TOKEN is missing in .env")

    application: Application = (
        ApplicationBuilder()
        .token(BOT_TOKEN)
        .connect_timeout(60)
        .read_timeout(120)
        .write_timeout(120)
        .pool_timeout(60)
        .media_write_timeout(300)
        .get_updates_connect_timeout(60)
        .get_updates_read_timeout(60)
        .get_updates_write_timeout(60)
        .get_updates_pool_timeout(60)
        .concurrent_updates(False)
        .build()
    )

    application.add_handler(CommandHandler("start", start))
    application.add_handler(CommandHandler("queue", queue_status))
    application.add_handler(CommandHandler("scan", scan))
    application.add_handler(CommandHandler("clear", clear))
    application.add_handler(CommandHandler("final", final))

    media_filter = (
        filters.PHOTO
        | filters.VIDEO
        | filters.ANIMATION
        | filters.Document.IMAGE
        | filters.Document.VIDEO
    )

    application.add_handler(MessageHandler(media_filter, collect_media))
    application.add_error_handler(error_handler)

    logger.info("Bot started.")
    application.run_polling(
        allowed_updates=Update.ALL_TYPES,
        drop_pending_updates=False,
    )


if __name__ == "__main__":
    main()