import { WebSocketServer as WSServer } from 'ws';
import dotenv from 'dotenv';

dotenv.config();

class MDWebSocketServer {
  constructor() {
    this.wss = null;
    this.connectedClients = new Map();
    this.currentJwtToken = null;
    this.currentCapacity = { internal: 0, external: 0 };
    this.instances = new Map();
    this.port = process.env.MESSAGE_DISPATCHER_WS_PORT || 8087;
  }

  async init() {
    return new Promise((resolve, reject) => {
      this.wss = new WSServer({ port: this.port });

      this.wss.on('listening', () => {
        console.log(`[MDWebSocketServer] WebSocket 服务已启动，端口: ${this.port}`);
        resolve();
      });

      this.wss.on('connection', (ws) => {
        console.log('[MDWebSocketServer] 新的 WebSocket 连接已建立');

        const clientId = Date.now().toString();
        this.connectedClients.set(clientId, ws);

        ws.on('message', (data) => {
          this.handleMessage(ws, data);
        });

        ws.on('close', () => {
          console.log('[MDWebSocketServer] WebSocket 连接已关闭');
          this.connectedClients.delete(clientId);
        });

        ws.on('error', (error) => {
          console.error('[MDWebSocketServer] WebSocket 连接错误:', error);
          this.connectedClients.delete(clientId);
        });
      });

      this.wss.on('error', (error) => {
        console.error('[MDWebSocketServer] WebSocket 服务错误:', error);
        reject(error);
      });
    });
  }

  handleMessage(ws, data) {
    try {
      const message = JSON.parse(data.toString());
      console.log(`[MDWebSocketServer] 收到消息: ${message.type}`);

      switch (message.type) {
        case 'JWT_UPDATE':
          this.handleJwtUpdate(message.data);
          break;
        case 'CAPACITY_UPDATE':
          this.handleCapacityUpdate(message.data);
          break;
        case 'INSTANCE_ONLINE':
          this.handleInstanceOnline(message.data);
          break;
        case 'INSTANCE_OFFLINE':
          this.handleInstanceOffline(message.data);
          break;
        case 'HEARTBEAT':
          this.handleHeartbeat(message.data, ws);
          break;
        default:
          console.log('[MDWebSocketServer] 未知消息类型:', message.type);
      }
    } catch (error) {
      console.error('[MDWebSocketServer] 解析消息失败:', error);
    }
  }

  handleJwtUpdate(data) {
    this.currentJwtToken = data.token;
    console.log('[MDWebSocketServer] JWT Token 已更新');
  }

  handleCapacityUpdate(data) {
    if (data.summary) {
      this.currentCapacity.internal = data.summary.onlineInstances - data.summary.busyInstances;
      console.log(`[MDWebSocketServer] 算力状态已更新: 内部可用 = ${this.currentCapacity.internal}`);
    }
  }

  handleInstanceOnline(data) {
    this.instances.set(data.instanceId, { ...data, status: 'online' });
    console.log(`[MDWebSocketServer] 实例上线: ${data.instanceId}`);
  }

  handleInstanceOffline(data) {
    this.instances.set(data.instanceId, { ...data, status: 'offline' });
    console.log(`[MDWebSocketServer] 实例下线: ${data.instanceId}`);
  }

  handleHeartbeat(data, ws) {
    ws.send(JSON.stringify({
      type: 'HEARTBEAT_ACK',
      data: { timestamp: new Date().toISOString() }
    }));
  }

  getJwtToken() {
    return this.currentJwtToken;
  }

  getInternalCapacity() {
    return this.currentCapacity.internal;
  }

  getExternalCapacity() {
    return parseInt(process.env.EXTERNAL_CAPACITY_MAX) || 10;
  }

  getInstances() {
    return Array.from(this.instances.values());
  }

  hasConnectedClients() {
    return this.connectedClients.size > 0;
  }
}

export default new MDWebSocketServer();