# -*- coding: utf-8 -*-
"""WebSocket 日志推送端点

提供 WS /ws/logs/{execution_id} 端点，实时推送 ETL 任务执行日志。
客户端连接后，先发送已有的历史日志行，再实时推送新日志，
直到执行结束（收到 None 哨兵）或客户端断开。

设计要点：
- 利用 TaskExecutor 已有的 subscribe/unsubscribe 机制
- 连接时先回放内存缓冲区中的历史日志，避免丢失已产生的行
- 通过 asyncio.Queue 接收实时日志，None 表示执行结束
"""

from __future__ import annotations

import logging
import time

from fastapi import APIRouter, WebSocket, WebSocketDisconnect

from ..services.task_executor import task_executor
# CHANGE 2026-03-24 | dev-trace-log: WebSocket 连接生命周期追踪
from ..trace.ws_wrapper import ws_trace_connect, ws_trace_disconnect, ws_trace_message

logger = logging.getLogger(__name__)

ws_router = APIRouter()


@ws_router.websocket("/ws/logs/{execution_id}")
async def ws_logs(websocket: WebSocket, execution_id: str) -> None:
    """实时推送指定 execution_id 的任务执行日志。

    流程：
    1. 接受 WebSocket 连接
    2. 回放内存缓冲区中已有的日志行
    3. 订阅 TaskExecutor，持续推送新日志
    4. 收到 None（执行结束）或客户端断开时关闭
    """
    await websocket.accept()
    logger.info("WebSocket 连接已建立: execution_id=%s", execution_id)

    # CHANGE 2026-03-24 | dev-trace-log: 记录 WS_CONNECT span
    _t0 = time.perf_counter()
    _ws_ctx = ws_trace_connect(execution_id, client_info=str(websocket.client))
    _msg_count = 0
    _total_bytes = 0

    # 订阅日志流
    queue = task_executor.subscribe(execution_id)
    _disconnect_reason = "normal"

    try:
        # 回放已有的历史日志行
        for line in task_executor.get_logs(execution_id):
            await websocket.send_text(line)
            _msg_count += 1
            _total_bytes += len(line.encode("utf-8"))
            ws_trace_message(_msg_count, _total_bytes)

        # 如果任务已经不在运行且没有订阅者队列中的数据，
        # 仍然保持连接等待——可能是任务刚结束但 queue 里还有未消费的消息
        while True:
            msg = await queue.get()
            if msg is None:
                # 执行结束哨兵
                break
            await websocket.send_text(msg)
            _msg_count += 1
            _total_bytes += len(msg.encode("utf-8"))
            ws_trace_message(_msg_count, _total_bytes)

    except WebSocketDisconnect:
        _disconnect_reason = "client_disconnect"
        logger.info("WebSocket 客户端断开: execution_id=%s", execution_id)
    except Exception:
        _disconnect_reason = "error"
        logger.exception("WebSocket 异常: execution_id=%s", execution_id)
    finally:
        task_executor.unsubscribe(execution_id, queue)
        # CHANGE 2026-03-24 | dev-trace-log: 记录 WS_DISCONNECT span 并写入日志
        _duration_ms = (time.perf_counter() - _t0) * 1000
        ws_trace_disconnect(_disconnect_reason, _msg_count, _duration_ms)
        # 安全关闭连接（客户端可能已断开，忽略错误）
        try:
            await websocket.close()
        except Exception:
            pass
        logger.info("WebSocket 连接已清理: execution_id=%s", execution_id)