ziper/bot.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.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", "200"))

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(
        """
        CREATE TABLE IF NOT EXISTS media_queue (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            chat_id INTEGER NOT NULL,
            user_id INTEGER NOT NULL,
            file_id TEXT NOT NULL,
            media_type TEXT NOT NULL,
            file_name TEXT,
            mime_type TEXT,
            created_at INTEGER NOT NULL
        )
        """
    )
    conn.commit()
    return conn


def add_to_queue(
    chat_id: int,
    user_id: int,
    file_id: str,
    media_type: str,
    file_name: Optional[str],
    mime_type: Optional[str],
) -> int:
    conn = db()
    cur = conn.cursor()
    cur.execute(
        """
        INSERT INTO media_queue
        (chat_id, user_id, file_id, media_type, file_name, mime_type, created_at)
        VALUES (?, ?, ?, ?, ?, ?, ?)
        """,
        (
            chat_id,
            user_id,
            file_id,
            media_type,
            file_name,
            mime_type,
            int(time.time()),
        ),
    )
    conn.commit()
    count = conn.execute(
        "SELECT COUNT(*) FROM media_queue WHERE chat_id = ? AND user_id = ?",
        (chat_id, user_id),
    ).fetchone()[0]
    conn.close()
    return count


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


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


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=92, 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",
                "2",
                str(output_path),
            ],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            timeout=120,
        )
        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 start(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    text = (
        "بات آماده است.\n\n"
        "مدیا بفرست: عکس، ویدیو، GIF، یا فایل image/video.\n\n"
        "دستورها:\n"
        "/queue - دیدن تعداد فایل‌های صف\n"
        "/final - دانلود همه مدیاها، تبدیل به عکس، ارسال خروجی\n"
        "/clear - پاک کردن صف"
    )
    await update.effective_message.reply_text(text)


async def queue_status(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
    chat_id = update.effective_chat.id
    user_id = update.effective_user.id
    rows = get_queue(chat_id, user_id)
    await update.effective_message.reply_text(f"تعداد مدیاهای داخل صف: {len(rows)}")


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 update.effective_message.reply_text(f"صف پاک شد. تعداد حذف‌شده: {deleted}")


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

    file_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
        media_type = "image"
        file_name = f"telegram_photo_{best_photo.file_unique_id}.jpg"
        mime_type = "image/jpeg"

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

    elif msg.animation:
        file_id = msg.animation.file_id
        media_type = "video"
        file_name = msg.animation.file_name or f"telegram_animation_{msg.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_name = doc.file_name or f"telegram_document_{doc.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:
            await msg.reply_text("این فایل image/video نیست؛ داخل صف اضافه نشد.")
            return

    else:
        await msg.reply_text("فقط عکس، ویدیو، GIF یا فایل image/video پشتیبانی می‌شود.")
        return

    count = add_to_queue(
        chat_id=chat_id,
        user_id=user_id,
        file_id=file_id,
        media_type=media_type,
        file_name=file_name,
        mime_type=mime_type,
    )

    await msg.reply_text(f"اضافه شد. تعداد فعلی صف: {count}\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 msg.reply_text("صف خالی است. اول چند مدیا بفرست.")
        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 msg.reply_text(f"شروع شد. تعداد فایل‌ها: {len(rows)}")

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

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

            await status.edit_text(f"درحال دانلود و تبدیل: {index}/{len(rows)}")

            tg_file = await context.bot.get_file(file_id)

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

            await tg_file.download_to_drive(custom_path=str(input_path))

            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)
            else:
                failed += 1

        if not output_images:
            await status.edit_text("هیچ خروجی تصویری ساخته نشد.")
            return

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

        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 context.bot.send_media_group(
                    chat_id=chat_id,
                    media=media_group,
                    reply_to_message_id=msg.message_id,
                )
            finally:
                for f in opened_files:
                    f.close()

        clear_queue(chat_id, user_id)

        done_text = f"تمام شد. خروجی‌های ارسال‌شده: {len(output_images)}"
        if failed:
            done_text += f"\nناموفق: {failed}"

        await status.edit_text(done_text)

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

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


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

    application: Application = ApplicationBuilder().token(BOT_TOKEN).build()

    application.add_handler(CommandHandler("start", start))
    application.add_handler(CommandHandler("queue", queue_status))
    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))

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


if __name__ == "__main__":
    main()