bot_handler.py

import logging
import requests
import json
import os
import time
import queue
import threading
import asyncio
from pyrogram import Client
from pyrogram.errors import FloodWait
from config import TELEGRAM_BOT_TOKEN, TELEGRAM_CHAT_ID, TELEGRAM_API_URL, API_ID, API_HASH, CACHE_MAX_AGE_MINUTES

logger = logging.getLogger(__name__)

CACHE_DIR = './cache'

# 全域快取進度追蹤與中斷控制
CACHE_PROGRESS = {}
CANCELED_TASKS = set()

# 大檔案 (Pyrogram) 專用排隊系統
PYRO_QUEUE = queue.Queue()
_pyro_worker_thread = None

def _pyro_worker():
    """專職處理大檔案的搬運工，一次只處理一個 Pyrogram 下載"""
    while True:
        try:
            file_id, app_instance, temp_path, cache_path, filename = PYRO_QUEUE.get()
            if file_id in CANCELED_TASKS:
                logger.info(f"[PyroWorker] Task {file_id} was canceled before starting.")
                PYRO_QUEUE.task_done(); continue
            
            logger.info(f"[PyroWorker] STARTING download for file: {filename} ({file_id})")
            try:
                # 執行實際的 Pyrogram 下載
                async def do_pyro():
                    session_name = f"c_{file_id[:8]}_{int(time.time())}"
                    app = Client(session_name, api_id=int(API_ID) if API_ID else None, api_hash=API_HASH, bot_token=TELEGRAM_BOT_TOKEN, in_memory=True)
                    await app.start()
                    try:
                        start_time = time.time()
                        def prog(current, total):
                            if file_id in CANCELED_TASKS: raise Exception("CANCELED")
                            elapsed = time.time() - start_time
                            speed = (current / 1024 / 1024) / elapsed if elapsed > 0 else 0
                            p = int((current / total) * 100) if total > 0 else 0
                            CACHE_PROGRESS[file_id].update({
                                'status': 'pyro_downloading', 
                                'percent': p, 
                                'downloaded': current, 
                                'total': total,
                                'speed': round(speed, 2)
                            })
                        await app.download_media(file_id, file_name=temp_path, progress=prog)
                        if os.path.exists(cache_path): os.remove(cache_path)
                        os.rename(temp_path, cache_path)
                        logger.info(f"[PyroWorker] SUCCESS: {filename} cached at {cache_path}")
                        CACHE_PROGRESS[file_id].update({'status': 'completed', 'percent': 100, 'speed': 0})
                    finally: await app.stop()
                
                asyncio.run(do_pyro())
            except FloodWait as e:
                logger.error(f"[PyroWorker] FLOOD WAIT triggered! Must wait for {e.value} seconds.")
                CACHE_PROGRESS[file_id].update({'status': 'failed', 'error': f"FloodWait: {e.value}s"})
            except Exception as e:
                if str(e) == "CANCELED": 
                    logger.info(f"[PyroWorker] CANCELED: {filename}")
                    CACHE_PROGRESS[file_id]['status'] = 'canceled'
                else: 
                    logger.error(f"[PyroWorker] FAILED for {filename}: {e}")
                    CACHE_PROGRESS[file_id].update({'status': 'failed', 'error': str(e)})
                if os.path.exists(temp_path): os.remove(temp_path)
            
            PYRO_QUEUE.task_done()
        except Exception as e: 
            logger.error(f"[PyroWorker] Critical loop error: {e}")
            time.sleep(1)

def start_pyro_worker():
    global _pyro_worker_thread
    if _pyro_worker_thread is None or not _pyro_worker_thread.is_alive():
        logger.info("[PyroWorker] Initializing background worker thread...")
        _pyro_worker_thread = threading.Thread(target=_pyro_worker, daemon=True)
        _pyro_worker_thread.start()

def cancel_cache_task(file_id):
    CANCELED_TASKS.add(file_id)
    if file_id in CACHE_PROGRESS: CACHE_PROGRESS[file_id]['status'] = 'canceled'

def remove_cache_task(file_id):
    if file_id in CACHE_PROGRESS: del CACHE_PROGRESS[file_id]
    if file_id in CANCELED_TASKS:
        try: CANCELED_TASKS.remove(file_id)
        except: pass

def download_file_to_cache(file_id, app=None, user_id=None):
    import database
    if not os.path.exists(CACHE_DIR): os.makedirs(CACHE_DIR, exist_ok=True)
    cache_file_path = os.path.join(CACHE_DIR, file_id)
    temp_cache_path = cache_file_path + ".tmp"
    
    if os.path.exists(cache_file_path) and (time.time() - os.path.getmtime(cache_file_path) < CACHE_MAX_AGE_MINUTES * 60):
        logger.info(f"[Cache] Hit for {file_id}, using existing file.")
        return True

    filename = file_id[:15]
    db_id = None
    if app:
        with app.app_context():
            file_obj = database.get_file_by_telegram_file_id(file_id, include_deleted=True)
            if file_obj:
                filename = file_obj.filename
                db_id = file_obj.id
    
    logger.info(f"[Cache] Initiating download for {filename}...")
    CACHE_PROGRESS[file_id] = {'status': 'starting', 'percent': 0, 'downloaded': 0, 'total': 0, 'filename': filename, 'id': db_id, 'user_id': user_id, 'speed': 0, 'error': None}

    url = f"https://api.telegram.org/bot{TELEGRAM_BOT_TOKEN}/getFile?file_id={file_id}"
    try:
        resp = requests.get(url); resp.raise_for_status()
        f_info = resp.json()['result']
        t_size = f_info.get('file_size', 0)
        telegram_download_url = f"https://api.telegram.org/file/bot{TELEGRAM_BOT_TOKEN}/{f_info['file_path']}"
        
        logger.info(f"[Cache] Small file path detected via Bot API. Starting direct stream...")
        start_time = time.time()
        with requests.get(telegram_download_url, stream=True) as r:
            r.raise_for_status()
            downloaded = 0
            with open(temp_cache_path, 'wb') as f:
                for chunk in r.iter_content(chunk_size=131072): # 加大緩衝區提升網速顯示穩定性
                    if file_id in CANCELED_TASKS:
                        f.close(); os.remove(temp_cache_path); return False
                    if not chunk: break
                    f.write(chunk); downloaded += len(chunk)
                    
                    elapsed = time.time() - start_time
                    speed = (downloaded / 1024 / 1024) / elapsed if elapsed > 0 else 0
                    p = int((downloaded / t_size) * 100) if t_size > 0 else 0
                    CACHE_PROGRESS[file_id].update({'status': 'downloading', 'percent': p, 'downloaded': downloaded, 'total': t_size, 'speed': round(speed, 2)})
        
        if os.path.exists(cache_file_path): os.remove(cache_file_path)
        os.rename(temp_cache_path, cache_file_path)
        logger.info(f"[Cache] Direct stream success: {filename}")
        CACHE_PROGRESS[file_id].update({'status': 'completed', 'percent': 100, 'speed': 0})
        return True

    except requests.exceptions.HTTPError as e:
        if e.response.status_code == 400 and "file is too big" in e.response.text:
            logger.info(f"[Cache] File too big for Bot API (>20MB). Switching to Pyrogram Worker queue...")
            start_pyro_worker()
            CACHE_PROGRESS[file_id]['status'] = 'waiting'
            PYRO_QUEUE.put((file_id, app, temp_cache_path, cache_file_path, filename))
            return True
        else:
            logger.error(f"[Cache] Bot API Error for {filename}: {e}")
            CACHE_PROGRESS[file_id].update({'status': 'failed', 'error': str(e)}); return False
    except Exception as e:
        logger.error(f"[Cache] Unexpected error for {filename}: {e}")
        CACHE_PROGRESS[file_id].update({'status': 'failed', 'error': str(e)})
        if os.path.exists(temp_cache_path): os.remove(temp_cache_path)
        return False

# (其餘管理函數、提取縮圖函數保持不變，略過以節省空間但實際文件中會保留...)
def get_current_cache_size_bytes():
    total_size = 0
    if os.path.exists(CACHE_DIR):
        for f in os.listdir(CACHE_DIR):
            fp = os.path.join(CACHE_DIR, f); total_size += os.path.getsize(fp) if os.path.isfile(fp) else 0
    return total_size

def _extract_official_thumbnail(item, result=None):
    if not item: return None
    if result and 'photo' in result and isinstance(result['photo'], list): return result['photo'][-1].get('file_id')
    if 'thumbs' in item and item['thumbs']: return max(item['thumbs'], key=lambda x: x.get('file_size', 0)).get('file_id')
    t = item.get('thumbnail') or item.get('thumb')
    return t.get('file_id') if (t and isinstance(t, dict)) else None

def clear_file_cache(file_id, thumbnail_id=None):
    for fid in [file_id, thumbnail_id]:
        if not fid: continue
        path = os.path.join(CACHE_DIR, fid)
        if os.path.exists(path):
            try: os.remove(path)
            except: pass

def clear_cache_manual():
    if os.path.exists(CACHE_DIR):
        for f in os.listdir(CACHE_DIR):
            try: os.remove(os.path.join(CACHE_DIR, f))
            except: pass

def clean_cache():
    from config import CACHE_MAX_AGE_MINUTES
    if not os.path.exists(CACHE_DIR): return
    now = time.time()
    for f in os.listdir(CACHE_DIR):
        path = os.path.join(CACHE_DIR, f)
        if os.path.getmtime(path) < now - (CACHE_MAX_AGE_MINUTES * 60):
            try: os.remove(path)
            except: pass

def upload_thumbnail_only(thumbnail_path):
    if not thumbnail_path or not os.path.exists(thumbnail_path): return None
    try:
        with open(thumbnail_path, 'rb') as f:
            r = requests.post(f"https://api.telegram.org/bot{TELEGRAM_BOT_TOKEN}/sendPhoto", files={'photo':('thumb.jpg', f, 'image/jpeg')}, data={'chat_id':TELEGRAM_CHAT_ID})
            return r.json()['result']['photo'][-1]['file_id'] if r.status_code == 200 else None
    except: return None

def send_file_to_telegram(file_storage, thumbnail_path=None, file_size=0):
    file_name = file_storage.filename; mime_type = file_storage.content_type
    method = "sendDocument"; files = {'document': (file_name, file_storage.stream, mime_type)}
    is_p = False
    if mime_type.startswith('video/'): method = "sendVideo"; files = {'video': (file_name, file_storage.stream, mime_type)}
    elif mime_type.startswith('audio/'): method = "sendAudio"; files = {'audio': (file_name, file_storage.stream, mime_type)}
    elif mime_type.startswith('image/'):
        if file_size < 10 * 1024 * 1024: method = "sendPhoto"; files = {'photo': (file_name, file_storage.stream, mime_type)}; is_p = True
        else: method = "sendDocument"; files = {'document': (file_name, file_storage.stream, mime_type)}
    o_t = None
    if not is_p and thumbnail_path and os.path.exists(thumbnail_path): o_t = open(thumbnail_path, 'rb'); files['thumb'] = o_t
    try:
        r = requests.post(f"https://api.telegram.org/bot{TELEGRAM_BOT_TOKEN}/{method}", files=files, data={'chat_id': TELEGRAM_CHAT_ID})
        rj = r.json()
        if rj.get('ok'):
            res = rj['result']
            item = res.get('video') or res.get('document') or res.get('audio') or (res.get('photo')[-1] if res.get('photo') else None)
            tid = _extract_official_thumbnail(item, result=res)
            if not tid and thumbnail_path and os.path.exists(thumbnail_path): tid = upload_thumbnail_only(thumbnail_path)
            return item.get('file_id'), file_name, mime_type, item.get('file_size'), res.get('message_id'), tid
    except Exception as e: logger.error(f"send_file error: {e}")
    finally:
        if o_t: o_t.close()
    return [None]*6

def send_media_group_to_telegram(tasks):
    url = f"https://api.telegram.org/bot{TELEGRAM_BOT_TOKEN}/sendMediaGroup"; o_f = []
    try:
        media = []; f_g = {}
        for i, task in enumerate(tasks):
            m_type = 'video' if task.mime_type.startswith('video/') else 'document'
            if task.mime_type.startswith('image/') and task.file_size < 5 * 1024 * 1024: m_type = 'photo'
            fs = open(task.temp_file_path, 'rb'); o_f.append(fs); f_g[f"f{i}"] = (task.filename, fs, task.mime_type)
            mi = {"type": m_type, "media": f"attach://f{i}"}
            if task.thumbnail_path and os.path.exists(task.thumbnail_path):
                tk = f"t{i}"; mi["thumb"] = f"attach://{tk}"; ft = open(task.thumbnail_path, 'rb'); o_f.append(ft); f_g[tk] = (os.path.basename(task.thumbnail_path), ft, 'image/jpeg')
            media.append(mi)
        r = requests.post(url, files=f_g, data={'chat_id': TELEGRAM_CHAT_ID, 'media': json.dumps(media)}); rj = r.json()
        if not rj.get('ok'): return None
        res = rj.get('result', []); ext = []
        for i, msg in enumerate(res):
            item = msg.get('video') or msg.get('document') or (msg.get('photo')[-1] if msg.get('photo') else None)
            if not item: continue
            tid = _extract_official_thumbnail(item, result=msg)
            if not tid and tasks[i].thumbnail_path and os.path.exists(tasks[i].thumbnail_path): tid = upload_thumbnail_only(tasks[i].thumbnail_path)
            ext.append((item['file_id'], item.get('file_name') or tasks[i].filename, item.get('mime_type') or tasks[i].mime_type, item.get('file_size'), msg.get('message_id'), tid, tasks[i]))
        return ext
    except Exception as e: logger.error(f"batch error: {e}"); return None
    finally:
        for f in o_f: f.close()

def stream_telegram_file_no_cache(file_id):
    try:
        r = requests.get(f"https://api.telegram.org/bot{TELEGRAM_BOT_TOKEN}/getFile?file_id={file_id}", timeout=10)
        if r.status_code == 200:
            path = r.json()['result']['file_path']
            with requests.get(f"https://api.telegram.org/file/bot{TELEGRAM_BOT_TOKEN}/{path}", stream=True) as rf:
                for c in rf.iter_content(chunk_size=8192): yield c
            return
    except: pass
    q = queue.Queue(); stop = threading.Event()
    def worker():
        async def stream_pyro():
            api_id_int = int(API_ID) if API_ID else None
            app_c = Client(f"s_{file_id[:8]}", api_id=api_id_int, api_hash=API_HASH, bot_token=TELEGRAM_BOT_TOKEN, in_memory=True)
            await app_c.start()
            try:
                async for c in app_c.stream_media(file_id):
                    if stop.is_set(): break
                    q.put(c)
            finally: await app_c.stop(); q.put(None)
        loop = asyncio.new_event_loop(); asyncio.set_event_loop(loop); loop.run_until_complete(stream_pyro()); loop.close()
    threading.Thread(target=worker, daemon=True).start()
    while True:
        c = q.get()
        if c is None: break
        yield c

def stream_and_cache_telegram_file(file_id, cancellable=True):
    cache_file_path = os.path.join(CACHE_DIR, file_id)
    if not os.path.exists(cache_file_path):
        if not download_file_to_cache(file_id): return None
    if os.path.exists(cache_file_path):
        def file_generator():
            with open(cache_file_path, 'rb') as f:
                while chunk := f.read(262144): yield chunk
        return file_generator()
    return None