services/progress_service.py

import asyncio
import logging
from typing import Dict, Callable, Optional

logger = logging.getLogger(__name__)

_progress_callbacks: Dict[str, list] = {}
_latest_progress: Dict[str, dict] = {}
_event_loop: Optional[asyncio.AbstractEventLoop] = None


def set_event_loop(loop: asyncio.AbstractEventLoop):
    global _event_loop
    _event_loop = loop


def register_progress_callback(video_id: str, callback: Callable):
    if video_id not in _progress_callbacks:
        _progress_callbacks[video_id] = []
    _progress_callbacks[video_id].append(callback)


def unregister_progress_callback(video_id: str, callback: Callable = None):
    if video_id in _progress_callbacks:
        if callback:
            _progress_callbacks[video_id] = [cb for cb in _progress_callbacks[video_id] if cb != callback]
        else:
            del _progress_callbacks[video_id]


def get_latest_progress(video_id: str) -> Optional[dict]:
    """获取最新的分析进度（用于 HTTP 轮询）"""
    return _latest_progress.get(video_id)


def set_progress_sync(video_id: str, stage: str, progress: float, message: str,
                      highlights_found: int = 0, eta_seconds: int = 0):
    data = {
        "stage": stage,
        "progress": round(progress, 2),
        "message": message,
        "highlights_found": highlights_found,
        "eta_seconds": eta_seconds,
    }
    _latest_progress[video_id] = data
    logger.info(f"[Progress] {video_id} | {stage} | {progress:.0%} | {message}")

    if video_id in _progress_callbacks and _event_loop and _event_loop.is_running():
        for callback in _progress_callbacks[video_id]:
            try:
                if asyncio.iscoroutinefunction(callback):
                    asyncio.run_coroutine_threadsafe(callback(data), _event_loop)
                else:
                    _event_loop.call_soon_threadsafe(callback, data)
            except Exception:
                pass


async def send_progress(video_id: str, stage: str, progress: float, message: str,
                        highlights_found: int = 0, eta_seconds: int = 0):
    """异步进度更新（从主线程调用，同时更新缓存和 WebSocket 回调）"""
    set_progress_sync(video_id, stage, progress, message, highlights_found, eta_seconds)

    if video_id in _progress_callbacks:
        for callback in _progress_callbacks[video_id]:
            try:
                if asyncio.iscoroutinefunction(callback):
                    await callback(_latest_progress[video_id])
                else:
                    callback(_latest_progress[video_id])
            except Exception:
                pass