diff --git a/bot.py b/bot.py new file mode 100644 index 0000000..bdf4612 --- /dev/null +++ b/bot.py @@ -0,0 +1,427 @@ +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()