/**
 * WebSocket 服务
 * 用于推送音频/视频生成完成事件
 */

import { Server as HttpServer } from 'http';
import WebSocket, { WebSocketServer } from 'ws';

const wss = new WebSocketServer({ noServer: true });

// 客户端连接管理
const clients = new Map<string, WebSocket>();

// ============ 客户端管理 ============

/**
 * 注册客户端连接
 */
export function addClient(clientId: string, ws: WebSocket) {
  clients.set(clientId, ws);
  console.log(`[WS] 客户端连接: ${clientId}, 当前在线: ${clients.size}`);
}

/**
 * 移除客户端连接
 */
export function removeClient(clientId: string) {
  clients.delete(clientId);
  console.log(`[WS] 客户端断开: ${clientId}, 当前在线: ${clients.size}`);
}

/**
 * 通过 clientId 发送消息
 */
export function sendToClient(clientId: string, event: string, data: any): boolean {
  const ws = clients.get(clientId);
  if (!ws || ws.readyState !== WebSocket.OPEN) {
    return false;
  }
  try {
    ws.send(JSON.stringify({ event, data }));
    return true;
  } catch (error) {
    console.error(`[WS] 发送消息失败: ${clientId}`, error);
    return false;
  }
}

/**
 * 广播消息到所有客户端
 */
export function broadcast(event: string, data: any) {
  const message = JSON.stringify({ event, data });
  clients.forEach((ws, clientId) => {
    if (ws.readyState === WebSocket.OPEN) {
      try {
        ws.send(message);
      } catch (error) {
        console.error(`[WS] 广播失败: ${clientId}`, error);
      }
    }
  });
}

// ============ 事件推送 ============

/**
 * 推送音频生成完成事件
 */
export function pushAudioGenerationComplete(bookId: string, chapterId: number, status: 'completed' | 'failed') {
  const event = 'audio_generation_complete';
  const data = { bookId, chapterId, status };
  console.log(`[WS] 推送 ${event}:`, data);
  // 广播给所有客户端（前端可根据 bookId 过滤）
  broadcast(event, data);
}

/**
 * 推送视频生成完成事件
 */
export function pushVideoGenerationComplete(bookId: string, chapterId: number, status: 'completed' | 'failed') {
  const event = 'video_generation_complete';
  const data = { bookId, chapterId, status };
  console.log(`[WS] 推送 ${event}:`, data);
  broadcast(event, data);
}

/**
 * 推送批量生成进度
 */
export function pushBatchGenerationProgress(taskId: string, step: string, progress: number) {
  const event = 'batch_generation_progress';
  const data = { taskId, step, progress };
  broadcast(event, data);
}

// ============ 初始化 ============

/**
 * 初始化 WebSocket 服务
 */
export function initWebSocket(server: HttpServer) {
  // 处理 HTTP upgrade 请求
  server.on('upgrade', (request, socket, head) => {
    const url = new URL(request.url || '', `http://${request.headers.host}`);

    // 只处理 /ws 路径的连接
    if (url.pathname !== '/ws') {
      socket.destroy();
      return;
    }

    wss.handleUpgrade(request, socket, head, (ws) => {
      // 从 query 参数获取 clientId
      const clientId = url.searchParams.get('clientId') || 'anonymous';
      addClient(clientId, ws);

      ws.on('close', () => {
        removeClient(clientId);
      });

      ws.on('error', (error) => {
        console.error(`[WS] 客户端错误: ${clientId}`, error);
        removeClient(clientId);
      });

      // 发送连接成功消息
      ws.send(JSON.stringify({ event: 'connected', data: { clientId } }));
    });
  });

  console.log('[WS] WebSocket 服务已初始化，挂载于 /ws 路径');
}

// 导出 wss 实例供外部使用
export { wss };