upload_worker.py

import threading
import time
import os
import logging
import asyncio
from pyrogram import Client

import database
import bot_handler
import user_handler
from models import db
from config import API_ID, API_HASH, TELEGRAM_BOT_TOKEN, TELEGRAM_CHAT_ID, USE_TG_OFFICIAL_THUMBNAIL
from pyrogram_clients import perform_pyrogram_upload

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("upload_worker")

# 鎖檔案路徑
LOCK_FILE = ".worker.lock"

class UploadWorker(threading.Thread):
    def __init__(self, app):
        super().__init__()
        self.app = app
        self.daemon = True
        self._running = True

    def run(self):
        logger.info("[WORKER-FLOW] Service initialized and running.")
        
        # --- 啟動時強力清理殘留鎖 ---
        if os.path.exists(LOCK_FILE):
            try:
                os.remove(LOCK_FILE)
                logger.info("[WORKER-FLOW] Stale lock file cleared on startup.")
            except: pass

        # --- 啟動時進行 Pyrogram 預熱 (Warmup) ---
        self._warmup_pyrogram()

        while self._running:
            try:
                # 1. 檢查並建立檔案鎖
                if os.path.exists(LOCK_FILE):
                    if time.time() - os.path.getmtime(LOCK_FILE) > 600:
                        try: os.remove(LOCK_FILE)
                        except: pass
                    else:
                        time.sleep(5); continue

                try:
                    with open(LOCK_FILE, "w") as f: f.write(str(os.getpid()))
                except: pass
                
                try:
                    with self.app.app_context():
                        self._process_tasks()
                        db.session.remove() 
                finally:
                    # 2. 釋放鎖 (增加安全檢查)
                    if os.path.exists(LOCK_FILE):
                        try: os.remove(LOCK_FILE)
                        except: pass
                    
            except Exception as e:
                logger.error(f"[WORKER-FLOW] Loop error: {e}")
                if os.path.exists(LOCK_FILE):
                    try: os.remove(LOCK_FILE)
                    except: pass
            time.sleep(5)

    def _warmup_pyrogram(self):
        """
        在背景嘗試獲取目標群組資訊，預熱 Pyrogram 並偵測賬號狀態（包括 FloodWait）。
        """
        from pyrogram.errors import FloodWait
        if not TELEGRAM_CHAT_ID:
            logger.warning("[WARMUP] TELEGRAM_CHAT_ID is not set, skipping warmup.")
            return

        logger.info(f"[WARMUP] Attempting to resolve TELEGRAM_CHAT_ID: {TELEGRAM_CHAT_ID}")
        
        if not os.path.exists("sessions"): os.makedirs("sessions", exist_ok=True)
        session_name = "sessions/shared_upload_session"

        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        try:
            async def do_warmup():
                c = Client(
                    session_name, 
                    api_id=int(API_ID) if API_ID else None, 
                    api_hash=API_HASH, 
                    bot_token=TELEGRAM_BOT_TOKEN,
                    workdir="."
                )
                try:
                    await c.start()
                    target_chat = int(TELEGRAM_CHAT_ID) if str(TELEGRAM_CHAT_ID).startswith('-100') else TELEGRAM_CHAT_ID
                    chat = await c.get_chat(target_chat)
                    logger.info(f"[WARMUP] Successfully resolved chat: {chat.title} ({chat.id})")
                    logger.info("[HEALTH-CHECK] Account status: OK (No active restrictions detected)")
                except FloodWait as e:
                    logger.error(f"[HEALTH-CHECK] Account is currently in FLOOD WAIT! Remaining: {e.value} seconds.")
                except Exception as e:
                    logger.error(f"[WARMUP] Failed to resolve chat {TELEGRAM_CHAT_ID}: {e}")
                finally:
                    if c.is_connected:
                        await c.stop()
            
            loop.run_until_complete(do_warmup())
        except Exception as e:
            logger.error(f"[WARMUP] Fatal error during warmup: {e}")
        finally:
            loop.close()

    def stop(self):
        self._running = False

    def _process_tasks(self):
        all_queued = database.get_upload_tasks(status='queued', limit=10)
        if not all_queued: return

        logger.info(f"[WORKER-FLOW] Found {len(all_queued)} queued tasks. Locking now.")
        
        all_ids = [t.id for t in all_queued]
        database.bulk_update_upload_task_status(all_ids, 'uploading')
        db.session.commit()
        db.session.expire_all() 
        
        tasks_to_process = []
        for tid in all_ids:
            t = database.get_upload_task_by_id(tid)
            if t: tasks_to_process.append(t)

        processed_ids = []
        
        # 0. 檢查超大檔案 (> 2000MB) 並分割
        for task in tasks_to_process:
            if task.file_size and task.file_size > 2000 * 1024 * 1024:
                logger.info(f"[LARGE-FILE-FLOW] Task {task.id} is too large ({task.file_size} bytes). Splitting...")
                self._split_large_file(task)
                processed_ids.append(task.id)
                continue

        media_group_candidates = [t for t in tasks_to_process if t.api_used == 'bot_api' and t.id not in processed_ids]
        
        # 1. 批次處理 (Media Group)
        if len(media_group_candidates) >= 2:
            batch = media_group_candidates[:10]
            batch_ids = [t.id for t in batch]
            logger.info(f"[BATCH-FLOW] Processing batch: {batch_ids}")
            for t in batch: self._ensure_thumbnail(t)
            self._process_media_group_batch(batch)
            processed_ids.extend(batch_ids)
        
        # 2. 單檔處理
        for task in tasks_to_process:
            if task.id in processed_ids: continue
            if task.status != 'uploading': continue
            
            logger.info(f"[SINGLE-FLOW] Processing Task {task.id}: {task.filename}")
            self._ensure_thumbnail(task)
            self._process_main_file_task(task)

    def _split_large_file(self, task):
        """
        將超大檔案分割成多個小於 2GB 的分片任務。
        媒體檔案 (MP4等) 用 ffmpeg 動態時長切片，非媒體檔案用 7z 分卷壓縮。
        """
        import subprocess
        import shutil
        import math
        
        file_path = task.temp_file_path
        if not os.path.exists(file_path):
            database.update_upload_task_status(task.id, 'failed', error_message="File not found for splitting.")
            db.session.commit(); return

        is_media = task.mime_type.startswith(('video/', 'audio/'))
        base_name, ext = os.path.splitext(task.filename)
        split_dir = os.path.join("./temp_uploads", f"split_{task.id}")
        os.makedirs(split_dir, exist_ok=True)
        
        parts = []
        try:
            if is_media:
                logger.info(f"[SPLIT-MEDIA] Analyzing duration for {task.filename}")
                # 1. 使用 ffprobe 獲取總時長
                probe_cmd = [
                    'ffprobe', '-v', 'error', '-show_entries', 'format=duration', 
                    '-of', 'default=noprint_wrappers=1:nokey=1', file_path
                ]
                try:
                    duration = float(subprocess.check_output(probe_cmd).decode('utf-8').strip())
                    # 2. 計算每份切片時間 (檔案大小 / 1.8GB = 份數; 時長 / 份數 = 每份秒數)
                    # 使用 1.8GB (1.8 * 1024 * 1024 * 1024 bytes) 作為安全閾值
                    safe_limit = 1.8 * 1024 * 1024 * 1024
                    num_parts = math.ceil(task.file_size / safe_limit)
                    segment_time = math.floor(duration / num_parts)
                    
                    logger.info(f"[SPLIT-MEDIA] Duration: {duration}s, Size: {task.file_size}, Parts: {num_parts}, Segment Time: {segment_time}s")
                except Exception as e:
                    logger.warning(f"[SPLIT-MEDIA] Failed to probe duration, fallback to 600s: {e}")
                    segment_time = 600 # 萬一獲取失敗，用保守的 10 分鐘

                # 3. 執行動態切片
                output_pattern = os.path.join(split_dir, f"{base_name}_part%03d{ext}")
                cmd = [
                    'ffmpeg', '-i', file_path, 
                    '-c', 'copy', '-map', '0', 
                    '-f', 'segment', 
                    '-segment_time', str(segment_time), 
                    '-reset_timestamps', '1',
                    output_pattern
                ]
                subprocess.run(cmd, check=True)
            else:
                logger.info(f"[SPLIT-NON-MEDIA] Using 7z for {task.filename}")
                # 非媒體檔案：使用 7z 進行物理分卷壓縮，每卷 1900M
                archive_path = os.path.join(split_dir, f"{task.filename}.7z")
                cmd = [
                    '7z', 'a', '-v1900m', archive_path, file_path
                ]
                subprocess.run(cmd, check=True)

            # 掃描生成的分片檔案
            for f in sorted(os.listdir(split_dir)):
                if f.endswith(ext) or '.7z.' in f:
                    parts.append(os.path.join(split_dir, f))

            if not parts:
                raise Exception("Splitting failed, no parts generated.")

            # 為每個分片創建新任務
            for i, part_path in enumerate(parts):
                part_size = os.path.getsize(part_path)
                part_filename = os.path.basename(part_path)
                api = 'pyrogram' if part_size >= 20 * 1024 * 1024 else 'bot_api'
                
                database.add_upload_task_to_db(
                    filename=part_filename,
                    temp_file_path=part_path,
                    target_folder=task.target_folder,
                    file_size=part_size,
                    mime_type=task.mime_type if is_media else 'application/x-7z-compressed',
                    api_used=api,
                    status='queued',
                    priority=task.priority + 1 
                )
            
            database.update_upload_task_status(task.id, 'completed', error_message=f"Split into {len(parts)} parts.")
            if os.path.exists(file_path): os.remove(file_path)
            logger.info(f"[SPLIT-FLOW] Task {task.id} split successfully into {len(parts)} parts.")
            
        except Exception as e:
            logger.error(f"[SPLIT-FLOW] Task {task.id} split failed: {e}")
            database.update_upload_task_status(task.id, 'failed', error_message=f"Split failed: {str(e)}")
        
        db.session.commit()
        
        db.session.commit()

    def _ensure_thumbnail(self, task):
        if USE_TG_OFFICIAL_THUMBNAIL: return
        if task.thumbnail_path and os.path.exists(task.thumbnail_path): return
        
        if task.mime_type.startswith(('image/', 'video/')):
            logger.info(f"[THUMB-FLOW] Custom mode: Generating thumbnail for Task {task.id}")
            
            # --- 關鍵修復：隔離循環 ---
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            try:
                path = loop.run_until_complete(user_handler.generate_thumbnail(task.temp_file_path, task.mime_type))
                if path: 
                    task.thumbnail_path = path
                    db.session.add(task)
                    db.session.commit()
            except Exception as e: logger.error(f"[THUMB-FLOW] Task {task.id} generation error: {e}")
            finally:
                loop.close()

    def _cleanup_successful_task(self, task):
        try:
            database.update_upload_task_status(task.id, 'completed')
            db.session.commit()
            if task.temp_file_path and os.path.exists(task.temp_file_path):
                try: os.remove(task.temp_file_path)
                except: pass
            if task.thumbnail_path and os.path.exists(task.thumbnail_path):
                try: os.remove(task.thumbnail_path)
                except: pass
            logger.info(f"[WORKER-FLOW] Task {task.id} finalized.")
        except Exception as e:
            logger.error(f"[WORKER-FLOW] Task {task.id} cleanup error: {e}")

    def _process_media_group_batch(self, tasks):
        try:
            extracted_data = bot_handler.send_media_group_to_telegram(tasks)
            if not extracted_data: raise Exception("MediaGroup returned NO DATA.")

            for data in extracted_data:
                f_id, f_name, m_type, size, msg_id, t_id, orig_task = data
                if not t_id and USE_TG_OFFICIAL_THUMBNAIL:
                    logger.info(f"[THUMB-FIX] Official thumbnail missing for {orig_task.filename}, triggering fallback.")
                    loop = asyncio.new_event_loop()
                    asyncio.set_event_loop(loop)
                    try:
                        temp_thumb = loop.run_until_complete(user_handler.generate_thumbnail(orig_task.temp_file_path, orig_task.mime_type))
                        if temp_thumb:
                            t_id = bot_handler.upload_thumbnail_only(temp_thumb)
                            if os.path.exists(temp_thumb): os.remove(temp_thumb)
                    except: pass
                    finally: loop.close()

                chat_str = str(TELEGRAM_CHAT_ID).replace('-100', '')
                database.add_file(f_name, f_id, orig_task.target_folder, size, m_type, t_id, f"https://t.me/c/{chat_str}/{msg_id}")
                self._cleanup_successful_task(orig_task)
        except Exception as e:
            logger.error(f"[BATCH-FLOW] Execution failed: {e}")
            t_ids = [t.id for t in tasks]
            database.bulk_update_upload_task_status(t_ids, 'failed', error_message=str(e))
            db.session.commit()

    def _process_main_file_task(self, task):
        try:
            thumb = task.thumbnail_path if task.thumbnail_path and os.path.exists(task.thumbnail_path) else None
            f_id = None; f_name = None; m_type = None; size = None; t_id = None; msg_id = None

            if task.api_used == 'pyrogram':
                # --- 關鍵修正：使用共享的持久 Session ---
                if not os.path.exists("sessions"): os.makedirs("sessions", exist_ok=True)
                session_name = "sessions/shared_upload_session"

                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)
                try:
                    async def run_pyro():
                        # 共享 Session 不需要每次刪除
                        c = Client(
                            session_name, 
                            api_id=int(API_ID) if API_ID else None, 
                            api_hash=API_HASH, 
                            bot_token=TELEGRAM_BOT_TOKEN,
                            workdir="."
                        )
                        await c.start()
                        try:
                            target_chat = int(TELEGRAM_CHAT_ID) if str(TELEGRAM_CHAT_ID).startswith('-100') else TELEGRAM_CHAT_ID
                            res = await perform_pyrogram_upload(c, task.temp_file_path, task.filename, target_chat, task.mime_type, task.file_size, thumbnail_path=thumb, task_id=task.id)
                            return res
                        finally:
                            await c.stop()
                    
                    f_id, f_name, m_type, size, t_id, msg_id = loop.run_until_complete(run_pyro())
                finally:
                    loop.close()
            else:
                from werkzeug.datastructures import FileStorage
                with open(task.temp_file_path, 'rb') as f:
                    fs = FileStorage(stream=f, filename=task.filename, content_type=task.mime_type)
                    f_id, f_name, m_type, size, msg_id, t_id = bot_handler.send_file_to_telegram(fs, thumbnail_path=thumb, file_size=task.file_size)

            if f_id:
                if not t_id and USE_TG_OFFICIAL_THUMBNAIL:
                    logger.info(f"[THUMB-FIX] Official thumbnail missing, triggering fallback.")
                    loop = asyncio.new_event_loop()
                    asyncio.set_event_loop(loop)
                    try:
                        temp_thumb = loop.run_until_complete(user_handler.generate_thumbnail(task.temp_file_path, task.mime_type))
                        if temp_thumb:
                            t_id = bot_handler.upload_thumbnail_only(temp_thumb)
                            if os.path.exists(temp_thumb): os.remove(temp_thumb)
                    except: pass
                    finally: loop.close()

                chat_str = str(TELEGRAM_CHAT_ID).replace('-100', '')
                database.add_file(f_name, f_id, task.target_folder, size, m_type, t_id, f"https://t.me/c/{chat_str}/{msg_id}")
                self._cleanup_successful_task(task)
            else:
                database.update_upload_task_status(task.id, 'failed', error_message="Upload failed - No File ID.")
                db.session.commit()
        except Exception as e:
            logger.error(f"[SINGLE-FLOW] Task {task.id} failed: {e}")
            database.update_upload_task_status(task.id, 'failed', error_message=str(e))
            db.session.commit()

upload_worker_thread = None
def start_upload_worker(app):
    global upload_worker_thread
    if upload_worker_thread is None or not upload_worker_thread.is_alive():
        upload_worker_thread = UploadWorker(app)
        upload_worker_thread.start()

def stop_upload_worker():
    global upload_worker_thread
    if upload_worker_thread is not None: upload_worker_thread.stop()