# --- routers/logs.py ---
import asyncio
import logging
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from logger import ColoredFormatter, register_websocket_handler

router = APIRouter()

class LogBroadcaster:
    """管理所有活躍 WebSocket 連線的廣播器 (Connection Manager)"""
    def __init__(self):
        self.active_connections: set[WebSocket] = set()
        self.loop: asyncio.AbstractEventLoop = None

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.add(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.discard(websocket)

    async def broadcast(self, message: str):
        """非同步推播日誌給所有連線中的客戶端"""
        if not self.active_connections:
            return
        
        # 併發發送，並透過 return_exceptions=True 確保單一連線異常不影響其他客戶端
        tasks = [self._safe_send(conn, message) for conn in self.active_connections]
        await asyncio.gather(*tasks, return_exceptions=True)

    async def _safe_send(self, websocket: WebSocket, message: str):
        try:
            await websocket.send_text(message)
        except Exception:
            self.disconnect(websocket)

# 建立全域廣播器實例
log_broadcaster = LogBroadcaster()

class WebSocketLogHandler(logging.Handler):
    """自訂 Log Handler，攔截系統日誌並安全地派發至非同步廣播器"""
    def __init__(self, broadcaster: LogBroadcaster):
        super().__init__()
        self.broadcaster = broadcaster
        # 沿用系統 ColoredFormatter，完美保留 ANSI 色碼
        self.setFormatter(ColoredFormatter())

    def emit(self, record):
        try:
            msg = self.format(record)
            loop = self.broadcaster.loop
            # 確保在 Event Loop 處於執行狀態時，安全地跨執行緒派發任務
            if loop and loop.is_running():
                loop.call_soon_threadsafe(
                    lambda: asyncio.create_task(self.broadcaster.broadcast(msg))
                )
        except Exception:
            self.handleError(record)

# 建立 Handler 實例並註冊至所有受管控的 Logger
websocket_log_handler = WebSocketLogHandler(log_broadcaster)
register_websocket_handler(websocket_log_handler)

@router.websocket("/ws/logs")
async def websocket_logs(websocket: WebSocket):
    """WebSocket 實時日誌串流端點"""
    # 若尚未綁定 Event Loop，則於首次連線時動態綁定
    if not log_broadcaster.loop:
        log_broadcaster.loop = asyncio.get_running_loop()

    await log_broadcaster.connect(websocket)
    try:
        # 保持連線，監聽客戶端斷線狀態
        while True:
            await websocket.receive_text()
    except WebSocketDisconnect:
        log_broadcaster.disconnect(websocket)
    except Exception:
        log_broadcaster.disconnect(websocket)