Spaces:

Thaimozhi
/

RenameBot

Running

File size: 10,520 Bytes

import os
import re
import shutil
import asyncio
import logging
import json

from pyrogram import Client, filters
from pyrogram.types import InlineKeyboardMarkup, InlineKeyboardButton, Message
from pyrogram.errors import MessageNotModified

logger = logging.getLogger(__name__)

# ─── /trim entry point ────────────────────────────────────────────────────────

@Client.on_message(filters.command("trim"))
async def trim_command(client, message: Message):
    """Must be a reply to a video/audio/document message."""
    if not message.reply_to_message:
        return await message.reply_text(
            "⚠️ **Reply to a video/audio file** with /trim to start.")

    reply = message.reply_to_message
    if not (reply.video or reply.audio or reply.document):
        return await message.reply_text("⚠️ Replied message must contain a video or audio file.")

    keyboard = InlineKeyboardMarkup([[
        InlineKeyboardButton("✂️ Split by Parts", callback_data=f"trim_parts_{reply.id}"),
        InlineKeyboardButton("⏱ Trim by Time",   callback_data=f"trim_time_{reply.id}"),
    ]])
    await message.reply_text(
        "🎬 **Trim / Split**\n\nChoose an operation:",
        reply_markup=keyboard, quote=True)


# ─── callback: split by parts ────────────────────────────────────────────────

@Client.on_callback_query(filters.regex(r"^trim_parts_"))
async def cb_trim_parts(client, query):
    m = re.match(r"^trim_parts_(\d+)$", query.data)
    if not m:
        return
    orig_msg_id = int(m.group(1))
    user_id = query.from_user.id
    await query.message.edit_text("🔢 **How many equal parts** do you want to split into?\n\n_Send a number (2–20):_")

    try:
        from pyromod.exceptions import ListenerTimeout
        resp = await client.listen(
            chat_id=query.message.chat.id, user_id=user_id, filters=filters.text, timeout=120)
        n_text = resp.text.strip()
        if not n_text.isdigit() or not (2 <= int(n_text) <= 20):
            await resp.delete()
            return await query.message.edit_text("❌ Please send a number between 2 and 20.")
        n_parts = int(n_text)
        await resp.delete()
    except ListenerTimeout:
        try:
            return await query.message.edit_text("⏳ Timed out. Send /trim again.")
        except MessageNotModified:
            return

    orig = await client.get_messages(query.message.chat.id, orig_msg_id)
    await _do_split_parts(client, query.message, orig, n_parts)


# ─── callback: trim by time ──────────────────────────────────────────────────

@Client.on_callback_query(filters.regex(r"^trim_time_"))
async def cb_trim_time(client, query):
    m = re.match(r"^trim_time_(\d+)$", query.data)
    if not m:
        return
    orig_msg_id = int(m.group(1))
    user_id = query.from_user.id
    await query.message.edit_text(
        "⏱ **Enter Start and End timestamps**\n\nFormat: `HH:MM:SS - HH:MM:SS`\nExample: `00:01:30 - 00:05:00`")

    try:
        from pyromod.exceptions import ListenerTimeout
        resp = await client.listen(
            chat_id=query.message.chat.id, user_id=user_id, filters=filters.text, timeout=120)
        ts_text = resp.text.strip()
        await resp.delete()
    except ListenerTimeout:
        try:
            return await query.message.edit_text("⏳ Timed out. Send /trim again.")
        except MessageNotModified:
            return

    # Support both "HH:MM:SS - HH:MM:SS" and "HH:MM:SS to HH:MM:SS"
    ts_text = ts_text.replace(" to ", " - ")
    sep_match = re.split(r'\s*-\s*', ts_text, maxsplit=1)
    if len(sep_match) != 2:
        return await query.message.edit_text(
            "❌ Invalid format. Use: `HH:MM:SS - HH:MM:SS`")

    start_ts, end_ts = sep_match[0].strip(), sep_match[1].strip()
    orig = await client.get_messages(query.message.chat.id, orig_msg_id)
    await _do_trim_time(client, query.message, orig, start_ts, end_ts)


# ─── helpers ─────────────────────────────────────────────────────────────────

async def _get_duration(file_path: str) -> float:
    """Return duration in seconds via ffprobe."""
    ffprobe = shutil.which('ffprobe')
    if not ffprobe:
        return 0.0
    cmd = [ffprobe, '-v', 'quiet', '-print_format', 'json',
           '-show_format', file_path]
    proc = await asyncio.create_subprocess_exec(
        *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
    stdout, _ = await proc.communicate()
    try:
        import json
        data = json.loads(stdout.decode())
        return float(data.get('format', {}).get('duration', 0))
    except Exception:
        return 0.0


async def _download_media_message(client, message, status_msg):
    """Download media from a message and return the local path."""
    if message.video:
        ext = ".mp4"
    elif message.audio:
        ext = f".{message.audio.mime_type.split('/')[-1]}" if message.audio.mime_type else ".mp3"
    elif message.document:
        name = message.document.file_name or "file"
        ext = os.path.splitext(name)[1] or ".mkv"
    else:
        ext = ".mp4"

    path = f"downloads/trim_{message.id}{ext}"
    os.makedirs("downloads", exist_ok=True)
    try:
        await status_msg.edit("**Downloading for processing...**")
    except MessageNotModified:
        pass
    file_path = await client.download_media(message, file_name=path)
    return file_path


async def _do_split_parts(client, status_msg, orig_msg, n_parts: int):
    ffmpeg = shutil.which('ffmpeg')
    if not ffmpeg:
        try:
            return await status_msg.edit("❌ FFmpeg not found on this server.")
        except MessageNotModified:
            return

    file_path = await _download_media_message(client, orig_msg, status_msg)
    if not file_path:
        try:
            return await status_msg.edit("❌ Could not download the file.")
        except MessageNotModified:
            return

    try:
        duration = await _get_duration(file_path)
        if duration <= 0:
            try:
                return await status_msg.edit("❌ Could not determine file duration.")
            except MessageNotModified:
                return

        part_duration = duration / n_parts
        base, ext = os.path.splitext(file_path)
        output_paths = []

        try:
            await status_msg.edit(f"✂️ Splitting into **{n_parts} parts**...")
        except MessageNotModified:
            pass

        for i in range(n_parts):
            start = i * part_duration
            out_path = f"{base}_part{i+1}{ext}"
            cmd = [
                ffmpeg, '-y',
                '-ss', str(start),
                '-i', file_path,
                '-t', str(part_duration),
                '-c', 'copy',
                '-loglevel', 'error',
                out_path
            ]
            proc = await asyncio.create_subprocess_exec(
                *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
            _, stderr = await proc.communicate()
            if proc.returncode != 0:
                logger.error(f"ffmpeg split error part {i+1}: {stderr.decode()}")
                await status_msg.edit(f"❌ FFmpeg error on part {i+1}: {stderr.decode()[:200]}")
                return
            output_paths.append(out_path)

        await status_msg.edit(f"📤 Uploading **{n_parts} parts**...")
        for i, out_path in enumerate(output_paths):
            part_name = f"Part {i+1} of {n_parts}{ext}"
            await client.send_document(
                chat_id=status_msg.chat.id,
                document=out_path,
                file_name=part_name,
                caption=f"✂️ Part {i+1}/{n_parts}"
            )
            os.remove(out_path)

        await status_msg.edit("✅ All parts uploaded!")
    finally:
        if os.path.exists(file_path):
            os.remove(file_path)


async def _do_trim_time(client, status_msg, orig_msg, start_ts: str, end_ts: str):
    ffmpeg = shutil.which('ffmpeg')
    if not ffmpeg:
        try:
            return await status_msg.edit("❌ FFmpeg not found on this server.")
        except MessageNotModified:
            return

    file_path = await _download_media_message(client, orig_msg, status_msg)
    if not file_path:
        try:
            return await status_msg.edit("❌ Could not download the file.")
        except MessageNotModified:
            return

    try:
        base, ext = os.path.splitext(file_path)
        out_path = f"{base}_trimmed{ext}"

        try:
            await status_msg.edit(f"✂️ Trimming `{start_ts}` → `{end_ts}`...")
        except MessageNotModified:
            pass

        cmd = [
            ffmpeg, '-y',
            '-ss', start_ts,
            '-to', end_ts,
            '-i', file_path,
            '-c', 'copy',
            '-loglevel', 'error',
            out_path
        ]
        proc = await asyncio.create_subprocess_exec(
            *cmd, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE)
        _, stderr = await proc.communicate()
        if proc.returncode != 0:
            return await status_msg.edit(f"❌ FFmpeg error: {stderr.decode()[:300]}")

        await status_msg.edit("📤 Uploading trimmed file...")

        trimmed_name = f"Trimmed_{start_ts.replace(':','-')}_{end_ts.replace(':','-')}{ext}"
        if orig_msg.video:
            await client.send_video(
                chat_id=status_msg.chat.id,
                video=out_path,
                file_name=trimmed_name,
                caption=f"✂️ Trimmed: `{start_ts}` → `{end_ts}`"
            )
        else:
            await client.send_document(
                chat_id=status_msg.chat.id,
                document=out_path,
                file_name=trimmed_name,
                caption=f"✂️ Trimmed: `{start_ts}` → `{end_ts}`"
            )

        if os.path.exists(out_path):
            os.remove(out_path)

        await status_msg.edit("✅ Trimmed file uploaded!")
    finally:
        if os.path.exists(file_path):
            os.remove(file_path)