File size: 19,920 Bytes
293f92d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
# app.py (Modified for Hugging Face + Telegram + Cloudflare)
import nest_asyncio
nest_asyncio.apply()

import os
import shutil
import asyncio
import traceback
import re
import time
from pathlib import Path

# Web Server & Bot
import httpx
from fastapi import FastAPI, BackgroundTasks
from pydantic import BaseModel
from contextlib import asynccontextmanager
import uvicorn

# Video Processing
try:
    from moviepy.editor import VideoFileClip
except ImportError:
    print("FATAL ERROR: moviepy is not installed. Run: pip install moviepy")
    # You might want to handle this more gracefully depending on HF deployment
    sys.exit(1)
import yt_dlp

# DNS Fix (Keep from reference)
import socket
import aiodns
original_getaddrinfo = socket.getaddrinfo
async def custom_getaddrinfo_async(host, port, family=0, type=0, proto=0, flags=0):
    resolver = aiodns.DNSResolver(nameservers=['8.8.8.8', '8.8.4.4'])
    try:
        result = await resolver.query(host, 'A'); addrlist = []
        for record in result: addrlist.append((socket.AF_INET, socket.SOCK_STREAM, 6, '', (record.host, port)))
        return addrlist
    except aiodns.error.DNSError:
        loop = asyncio.get_running_loop(); return await loop.run_in_executor(None, original_getaddrinfo, host, port, family, type, proto, flags)
def custom_getaddrinfo_sync(*args, **kwargs):
    try: loop = asyncio.get_running_loop()
    except RuntimeError: loop = asyncio.new_event_loop(); asyncio.set_event_loop(loop)
    return loop.run_until_complete(custom_getaddrinfo_async(*args, **kwargs))
socket.getaddrinfo = custom_getaddrinfo_sync

# --- CONFIGURATION ---
WORKER_URL = os.environ.get("WORKER_URL") # Get Worker URL from HF Secrets
LOCAL_TEMP_FOLDER = "temp_processing"
# Ensure the temp folder exists at startup
os.makedirs(LOCAL_TEMP_FOLDER, exist_ok=True)

# User session storage
USER_SESSIONS = {}

# --- TELEGRAM UTILS (Assumed to exist in utils/telegram_utils.py or similar) ---
# You need functions that POST to your WORKER_URL
# Example (replace with your actual implementation):
async def send_telegram_message(chat_id, text, reply_markup=None):
    url = f"{WORKER_URL}"
    payload = {'action': 'sendMessage', 'chat_id': str(chat_id), 'text': text}
    if reply_markup:
        payload['reply_markup'] = json.dumps(reply_markup)
    try:
        async with httpx.AsyncClient() as client:
            response = await client.post(url, json=payload, timeout=60)
            response.raise_for_status()
            print(f"Sent message to {chat_id}: {response.json()}")
    except Exception as e:
        print(f"ERROR sending message to {chat_id}: {e}")

async def send_telegram_video(chat_id, video_path, caption):
    url = f"{WORKER_URL}"
    with open(video_path, 'rb') as video_file:
        files = { 'video': (os.path.basename(video_path), video_file, 'video/mp4') }
        data = { 'action': 'sendVideo', 'chat_id': str(chat_id), 'caption': caption }
        try:
            async with httpx.AsyncClient() as client:
                # Increased timeout for potentially large files
                response = await client.post(url, data=data, files=files, timeout=600)
                response.raise_for_status()
                print(f"Sent video to {chat_id}: {response.json()}")
                return True # Indicate success
        except Exception as e:
            print(f"ERROR sending video to {chat_id}: {e}")
            await send_telegram_message(chat_id, f"❌ Lỗi khi gửi file video: {e}")
            return False # Indicate failure

async def answer_telegram_callback_query(callback_query_id, text="OK"):
    url = f"{WORKER_URL}"
    payload = {'action': 'answerCallbackQuery', 'callback_query_id': callback_query_id, 'text': text}
    try:
        async with httpx.AsyncClient() as client:
            response = await client.post(url, json=payload, timeout=30)
            response.raise_for_status()
    except Exception as e:
        print(f"ERROR answering callback query {callback_query_id}: {e}")

async def edit_telegram_message_text(chat_id, message_id, text, reply_markup=None):
    url = f"{WORKER_URL}"
    payload = {'action': 'editMessageText', 'chat_id': str(chat_id), 'message_id': message_id, 'text': text}
    if reply_markup:
        payload['reply_markup'] = json.dumps(reply_markup)
    try:
        async with httpx.AsyncClient() as client:
            response = await client.post(url, json=payload, timeout=60)
            response.raise_for_status()
    except Exception as e:
        print(f"ERROR editing message {message_id} in chat {chat_id}: {e}")


# --- VIDEO PROCESSING FUNCTIONS (Adapted from GUI code) ---

# Helper to get script directory (important for finding cookies.txt)
def get_script_dir():
    # In HF Spaces, the current working directory is usually the repo root
    return os.getcwd()

async def mute_video_with_moviepy_hf(original_video_path, chat_id):
    """ Mutes video using MoviePy, sends status messages via Telegram. """
    await send_telegram_message(chat_id, f"   Bắt đầu tắt tiếng: {os.path.basename(original_video_path)}")
    original_path = Path(original_video_path)
    muted_filename = original_path.parent / f"{original_path.stem}_MUTED.mp4"
    video_clip = None
    loop = asyncio.get_running_loop()

    try:
        def sync_mute_copy():
            nonlocal video_clip # Allow modifying outer scope variable
            video_clip = VideoFileClip(str(original_path))
            final_clip = video_clip.set_audio(None)
            final_clip.write_videofile(str(muted_filename), codec='copy', logger=None)

        await send_telegram_message(chat_id, "      Đang thử tắt tiếng (copy - nhanh)...")
        await loop.run_in_executor(None, sync_mute_copy) # Run synchronous moviepy code in executor
        await send_telegram_message(chat_id, f"      Tắt tiếng (copy) thành công: {muted_filename.name}")

    except Exception as e:
        await send_telegram_message(chat_id, f"      Lỗi 'codec=copy': {str(e)}. Đang thử nén lại (chậm hơn)...")
        try:
            def sync_mute_reencode():
                # Need to reopen the clip within the executor context
                clip_retry = VideoFileClip(str(original_path))
                final_clip_retry = clip_retry.set_audio(None)
                final_clip_retry.write_videofile(
                    str(muted_filename), threads=4, codec="libx264", audio_codec="aac", logger=None
                )
                clip_retry.close() # Close the clip explicitly

            # Close the original clip if it exists before retrying
            if video_clip: video_clip.close(); video_clip = None

            await loop.run_in_executor(None, sync_mute_reencode) # Run re-encode in executor
            await send_telegram_message(chat_id, f"      Tắt tiếng (nén lại) thành công: {muted_filename.name}")
            return (True, str(muted_filename))
        except Exception as e2:
            # Explicitly close clip on re-encode error if it's still somehow open
            # (though sync_mute_reencode should handle it)
            if 'clip_retry' in locals() and locals()['clip_retry']: locals()['clip_retry'].close()

            await send_telegram_message(chat_id, f"      LỖI MOVIEPY (lần 2):\n{str(e2)}")
            return (False, f"LỖI MOVIEPY:\n{str(e2)}")
    finally:
        if video_clip: video_clip.close()
    return (True, str(muted_filename))

async def download_single_video_hf(url, output_path, chat_id):
    """ Downloads a single video using yt-dlp, sends status via Telegram.

        Uses COOKIE_STRING secret instead of cookies.txt file. """
    await send_telegram_message(chat_id, f"Bắt đầu tải: {url}")
    original_filename = None
    loop = asyncio.get_running_loop()

    # === THAY ĐỔI: Lấy cookie từ Secret ===
    cookie_string = os.environ.get("COOKIE_STRING")
    if not cookie_string:
        await send_telegram_message(chat_id, "❌ LỖI NGHIÊM TRỌNG: Secret 'COOKIE_STRING' chưa được thiết lập trên Hugging Face!")
        return (False, "Lỗi Server: Thiếu Secret Cookie", None)
    # =======================================

    # Giữ nguyên Headers khác, nhưng thêm header 'Cookie'
    YTDLP_HEADERS = {
        'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36',
        'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7',
        'Accept-Language': 'en-US,en;q=0.9,vi;q=0.8',
        # === THAY ĐỔI: Thêm header Cookie ===
        'Cookie': cookie_string
        # =======================================
    }
    output_template = os.path.join(output_path, '%(title)s - %(id)s.%(ext)s')
    downloaded_file_holder = []

    def progress_hook(d):
        if d['status'] == 'finished':
            filepath = d.get('filename') or d.get('info_dict', {}).get('_filename')
            if filepath and filepath not in downloaded_file_holder:
                 downloaded_file_holder.append(filepath)

    ydl_opts = {
        'format': 'best',
        'outtmpl': output_template,
        'merge_output_format': 'mp4',
        'noplaylist': True,
        # 'cookiefile': cookie_file_path, # <<<=== XÓA DÒNG NÀY
        'http_headers': YTDLP_HEADERS, # Headers giờ đã chứa cookie
        'quiet': True, 'verbose': False, 'no_warnings': True,
        'ignoreerrors': False,
        'progress_hooks': [progress_hook],
        'paths': {'temp': os.path.join(LOCAL_TEMP_FOLDER, 'ytdlp_temp')},
        'fragment_retries': 10,
        'retries': 10,
    }

    try:
        def sync_download():
             with yt_dlp.YoutubeDL(ydl_opts) as ydl:
                 ydl.download([url])

        await loop.run_in_executor(None, sync_download)

        if not downloaded_file_holder:
             raise yt_dlp.utils.DownloadError("Không xác định được tên file đã tải.")
        original_filename = downloaded_file_holder[0]
        await send_telegram_message(chat_id, f"   ✅ Tải gốc xong: {os.path.basename(original_filename)}")
        return (True, "Tải gốc OK", original_filename)

    except Exception as e:
        error_message = str(e)
        await send_telegram_message(chat_id, f"   ❌ LỖI KHI TẢI GỐC:\n{error_message}")
        # Specific error checking
        # === THAY ĐỔI: Kiểm tra lỗi cookie kỹ hơn ===
        if "HTTP Error 403" in error_message or "401 Unauthorized" in error_message or "Fresh cookies" in error_message:
            return (False, "Lỗi: Cookies hết hạn hoặc không hợp lệ (kiểm tra Secret).", None)
        # ===========================================
        if "Unsupported URL" in error_message: return (False, "Lỗi: Link không hỗ trợ.", None)
        if "Private video" in error_message: return (False, "Lỗi: Video riêng tư.", None)
        if "Video unavailable" in error_message: return (False, "Lỗi: Video không khả dụng.", None)
        # General error
        return (False, f"Lỗi tải: {error_message[:100]}...", None)

# --- Background Task for Processing ---
async def process_links_task(chat_id, urls, is_muting):
    """ Background task to download and optionally mute multiple videos. """
    success_count = 0
    fail_count = 0
    mute_fail_count = 0
    total_links = len(urls)
    processed_files_info = [] # Store tuples of (original_path, muted_path_or_error)

    # Create a unique temp folder for this task
    task_temp_folder = os.path.join(LOCAL_TEMP_FOLDER, f"task_{chat_id}_{int(time.time())}")
    os.makedirs(task_temp_folder, exist_ok=True)

    await send_telegram_message(chat_id, f"🚀 Bắt đầu xử lý {total_links} link...")

    for i, url in enumerate(urls):
        await send_telegram_message(chat_id, f"--- Xử lý link {i+1}/{total_links} ---")
        original_filepath = None
        muted_filepath_or_error = None
        try:
            # Download to the task's temp folder
            dl_success, dl_msg, original_filepath = await download_single_video_hf(url, task_temp_folder, chat_id)

            if dl_success:
                if is_muting:
                    mute_success, mute_result = await mute_video_with_moviepy_hf(original_filepath, chat_id)
                    if mute_success:
                        muted_filepath_or_error = mute_result
                    else:
                        mute_fail_count += 1
                        muted_filepath_or_error = f"Lỗi mute: {mute_result}" # Store mute error
                        # Still counts as a download success
                success_count += 1
                processed_files_info.append((original_filepath, muted_filepath_or_error))
            else:
                fail_count += 1
                # dl_msg contains the error message from download_single_video_hf
                await send_telegram_message(chat_id, f"⚠️ Lỗi với link {i+1}: {dl_msg}")

        except Exception as e:
             fail_count += 1
             await send_telegram_message(chat_id, f"❌ Lỗi hệ thống nghiêm trọng khi xử lý link {i+1}: {e}")
             traceback.print_exc() # Print full traceback to server logs

        # Optional: Short delay between downloads
        await asyncio.sleep(1)

    # --- Sending results back ---
    await send_telegram_message(chat_id, f"--- Hoàn tất tải {success_count}/{total_links} video gốc ---")
    if is_muting:
         await send_telegram_message(chat_id, f"--- Bắt đầu gửi {success_count} video (bản gốc hoặc đã tắt tiếng) ---")

    sent_count = 0
    send_errors = 0
    for i, (orig_path, mute_path_or_err) in enumerate(processed_files_info):
        file_to_send = None
        caption = os.path.basename(orig_path or "video")

        if is_muting:
            if isinstance(mute_path_or_err, str) and not mute_path_or_err.startswith("Lỗi mute"):
                file_to_send = mute_path_or_err # Send muted version if success
                caption = os.path.basename(mute_path_or_err) + " (Đã tắt tiếng)"
            else:
                file_to_send = orig_path # Send original if muting failed
                caption += f" (Lỗi tắt tiếng: {mute_path_or_err})"
        else:
            file_to_send = orig_path # Send original if not muting

        if file_to_send and os.path.exists(file_to_send):
            await send_telegram_message(chat_id, f"   Đang gửi video {i+1}/{len(processed_files_info)}...")
            if await send_telegram_video(chat_id, file_to_send, caption):
                sent_count += 1
            else:
                send_errors += 1
        else:
            # This case means download failed, already logged earlier
            pass

    # --- Final Summary ---
    summary_message = f"🏁 Hoàn thành!\n- Tải thành công: {success_count}/{total_links}"
    if fail_count > 0: summary_message += f"\n- Tải thất bại: {fail_count}"
    if is_muting:
        if mute_fail_count > 0: summary_message += f"\n- Lỗi tắt tiếng: {mute_fail_count}"
    if send_errors > 0: summary_message += f"\n- Lỗi gửi file: {send_errors}"

    await send_telegram_message(chat_id, summary_message)

    # --- Cleanup ---
    if chat_id in USER_SESSIONS: del USER_SESSIONS[chat_id]
    try:
        shutil.rmtree(task_temp_folder)
        print(f"Cleaned up temp folder: {task_temp_folder}")
    except Exception as e:
        print(f"ERROR cleaning up temp folder {task_temp_folder}: {e}")


# --- FastAPI Setup ---
@asynccontextmanager
async def lifespan(app: FastAPI):
    print("🚀 Server is starting...")
    # Initial cleanup just in case
    if os.path.exists(LOCAL_TEMP_FOLDER): shutil.rmtree(LOCAL_TEMP_FOLDER)
    os.makedirs(LOCAL_TEMP_FOLDER, exist_ok=True)
    print("✅ Bot ready.")
    yield
    print("👋 Server shutting down.")

app = FastAPI(lifespan=lifespan)

class TelegramUpdate(BaseModel):
    message: dict | None = None
    callback_query: dict | None = None

@app.post("/webhook")
async def handle_webhook(update: TelegramUpdate, background_tasks: BackgroundTasks):
    if update.callback_query:
        # --- Handle Button Clicks ---
        callback_query = update.callback_query
        chat_id = callback_query["message"]["chat"]["id"]
        message_id = callback_query["message"]["message_id"]
        callback_data = callback_query["data"]
        callback_id = callback_query["id"]

        session = USER_SESSIONS.get(chat_id)
        if not session or session.get("step") != "awaiting_mute_choice":
            await answer_telegram_callback_query(callback_id, "⚠️ Lựa chọn đã hết hạn hoặc không hợp lệ. Gửi lại link nhé.")
            return {"status": "ok, session invalid"}

        # Remove the inline keyboard from the original message
        await edit_telegram_message_text(chat_id, message_id, callback_query["message"]["text"], reply_markup={})
        await answer_telegram_callback_query(callback_id, "OK, bắt đầu xử lý!") # Acknowledge button press

        # Store mute choice and start background task
        is_muting = (callback_data == "mute_yes")
        session["is_muting"] = is_muting
        session["step"] = "processing"

        background_tasks.add_task(process_links_task, chat_id, session["urls"], is_muting)
        return {"status": "ok, processing started"}

    elif update.message:
        # --- Handle Incoming Messages ---
        message = update.message
        chat_id = message["chat"]["id"]
        text = message.get("text", "").strip()

        if text == "/start":
            if chat_id in USER_SESSIONS: del USER_SESSIONS[chat_id]
            await send_telegram_message(chat_id, "👋 Chào bạn! Gửi 1 hoặc nhiều link video (Douyin/XHS), mỗi link một dòng.")
            return {"status": "ok, start message"}

        # Check if currently processing for this user
        if USER_SESSIONS.get(chat_id, {}).get("step") == "processing":
             await send_telegram_message(chat_id, "⏳ Bot đang xử lý yêu cầu trước đó của bạn. Vui lòng chờ nhé.")
             return {"status": "ok, busy"}

        # Extract URLs from the message
        url_pattern = r'https?://[^\s]+'
        urls = [match.group(0) for match in re.finditer(url_pattern, text)]

        if not urls:
            await send_telegram_message(chat_id, "⚠️ Không tìm thấy link hợp lệ nào. Vui lòng gửi link video.")
            return {"status": "ok, no links found"}

        # Store URLs and ask about muting
        USER_SESSIONS[chat_id] = {
            "step": "awaiting_mute_choice",
            "urls": urls
        }
        keyboard = { "inline_keyboard": [
            [{"text": "🔇 Có, tắt tiếng", "callback_data": "mute_yes"}],
            [{"text": "🔊 Không, giữ nguyên", "callback_data": "mute_no"}]
        ]}
        await send_telegram_message(chat_id, f"Tìm thấy {len(urls)} link. Bạn có muốn tạo thêm bản *tắt tiếng* cho video không?", reply_markup=keyboard)
        return {"status": "ok, links received, asking mute"}

    return {"status": "ok, ignored"} # Ignore other update types


@app.get("/")
def read_root(): return {"message": "Downloader Bot is running."}

# --- To run locally (optional) ---
# if __name__ == "__main__":
#    uvicorn.run(app, host="0.0.0.0", port=7860)