import { WebSocketServer, WebSocket } from 'ws';

// 客户端连接管理
const clients = new Map<string, WebSocket[]>(); // merchantId -> ws[]
let wsServer: WebSocketServer | null = null;

// 注册客户端连接
export const registerClient = (merchantId: string, ws: WebSocket): void => {
  if (!clients.has(merchantId)) {
    clients.set(merchantId, []);
  }
  clients.get(merchantId)?.push(ws);
  console.log(`[WebSocket] Client registered for merchant: ${merchantId}`);
};

// 移除客户端连接
export const removeClient = (merchantId: string, ws: WebSocket): void => {
  if (clients.has(merchantId)) {
    const clientList = clients.get(merchantId);
    if (clientList) {
      const index = clientList.indexOf(ws);
      if (index > -1) {
        clientList.splice(index, 1);
        console.log(`[WebSocket] Client removed for merchant: ${merchantId}`);
      }
      
      // 如果该商户没有客户端连接了，移除该商户
      if (clientList.length === 0) {
        clients.delete(merchantId);
        console.log(`[WebSocket] Merchant removed: ${merchantId}`);
      }
    }
  }
};

// 向特定商户推送消息
export const pushToMerchant = (merchantId: string, data: any): void => {
  const clientList = clients.get(merchantId);
  if (clientList) {
    const message = JSON.stringify(data);
    clientList.forEach(ws => {
      if (ws.readyState === WebSocket.OPEN) {
        ws.send(message);
      }
    });
    console.log(`[WebSocket] Pushed message to merchant: ${merchantId}`);
  }
};

// 向所有客户端推送消息
export const broadcastToAll = (data: any): void => {
  const message = JSON.stringify(data);
  clients.forEach((clientList) => {
    clientList.forEach(ws => {
      if (ws.readyState === WebSocket.OPEN) {
        ws.send(message);
      }
    });
  });
  console.log('[WebSocket] Broadcasted message to all clients');
};

// 启动WebSocket服务器
export const startWebSocketServer = (): void => {
  // 创建WebSocket服务器
  wsServer = new WebSocketServer({ port: 8085 });
  
  wsServer.on('connection', (ws, req) => {
    // 从请求中获取merchantId（实际应用中应该从认证信息中获取）
    const url = new URL(req.url || '', 'http://localhost');
    const merchantId = url.searchParams.get('merchantId') || 'anonymous';
    
    registerClient(merchantId, ws);
    
    // 发送连接成功消息
    ws.send(JSON.stringify({
      type: 'CONNECTION_SUCCESS',
      message: 'Connected to WebSocket server',
      merchantId
    }));
    
    // 处理客户端消息
    ws.on('message', (message) => {
      console.log(`[WebSocket] Received message from ${merchantId}:`, message.toString());
    });
    
    // 处理连接关闭
    ws.on('close', () => {
      removeClient(merchantId, ws);
    });
    
    // 处理连接错误
    ws.on('error', (error) => {
      console.error(`[WebSocket] Error for ${merchantId}:`, error);
      removeClient(merchantId, ws);
    });
  });
  
  console.log('[WebSocket] Server started on port 8085');
};

export default {
  startWebSocketServer,
  pushToMerchant,
  broadcastToAll,
  registerClient,
  removeClient
};
