import { WebSocketServer as WSServer } from 'ws';
import logger from '../logger/index.js';
import bridgeManager from '../bridge-manager/index.js';
import taskScheduler from '../task-scheduler/index.js';
import mdWebSocketClient from '../md-websocket-client/index.js';
import { v4 as uuidv4 } from 'uuid';

class WebSocketServer {
  constructor() {
    this.wss = null;
    this.pendingRequests = new Map();
    this.instanceTaskMap = new Map();
    this.TASK_TIMEOUT = 5 * 60 * 1000;
  }

  start(server) {
    this.wss = new WSServer({ 
      server,
      keepalive: true
    });
    logger.info('WebSocket 服务器已启动');

    this.wss.on('connection', (ws) => {
      this.handleConnection(ws);
    });
  }

  handleConnection(ws) {
    let bridgeId = null;
    let pingInterval = null;
    let pongTimeout = null;

    logger.info('新的 WebSocket 连接已建立');

    const PING_INTERVAL = 30000;
    const PONG_TIMEOUT = 10000;

    const sendPing = () => {
      if (ws.readyState !== WSServer.OPEN) {
        clearInterval(pingInterval);
        return;
      }
      ws.ping();
      pongTimeout = setTimeout(() => {
        logger.warn('PONG 响应超时，关闭连接');
        ws.terminate();
      }, PONG_TIMEOUT);
    };

    pingInterval = setInterval(sendPing, PING_INTERVAL);

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

    ws.on('pong', () => {
      if (pongTimeout) {
        clearTimeout(pongTimeout);
        pongTimeout = null;
      }
    });

    ws.on('close', (code, reason) => {
      clearInterval(pingInterval);
      if (pongTimeout) {
        clearTimeout(pongTimeout);
      }
      if (bridgeId) {
        this.handleBridgeDisconnect(bridgeId);
        bridgeManager.unregisterBridge(bridgeId);
        this.cleanupPendingRequests(bridgeId);
      }
      logger.info(`WebSocket 连接已关闭 (code: ${code})`);
    });

    ws.on('error', (error) => {
      clearInterval(pingInterval);
      if (pongTimeout) {
        clearTimeout(pongTimeout);
      }
      logger.error('WebSocket 连接错误:', error);
    });
  }

  async handleBridgeDisconnect(bridgeId) {
    const bridge = bridgeManager.getBridge(bridgeId);
    if (bridge?.info?.instances) {
      for (const instance of bridge.info.instances) {
        const affectedTaskId = bridgeManager.handleInstanceOffline(instance.id);
        if (affectedTaskId) {
          taskScheduler.recoverTask(affectedTaskId, 'bridge_disconnect');
        }
      }
    }
    const capacity = bridgeManager.getAvailableCapacity();
    await taskScheduler.setCurrentCapacity(capacity.online);
  }

  async handleMessage(ws, data, setBridgeId) {
    try {
      const message = JSON.parse(data.toString());
      logger.debug(`收到消息：${message.type}`);

      switch (message.type) {
        case 'REGISTER':
          this.handleRegister(ws, message, setBridgeId);
          break;
        case 'HEARTBEAT':
          this.handleHeartbeat(message);
          break;
        case 'TASK_ACK':
          this.handleTaskAck(message);
          break;
        case 'TASK_END':
          this.handleTaskEnd(message);
          break;
        case 'INSTANCE_CHECK_ACK':
          this.handleBridgeResponse(message);
          break;
        case 'INSTANCE_STATUS_UPDATE':
          this.handleInstanceStatusUpdate(message);
          break;
        case 'PONG':
          break;
        default:
          logger.debug('未知消息类型:', message.type);
      }
    } catch (error) {
      logger.error('解析消息失败:', error);
    }
  }

  async handleRegister(ws, message, setBridgeId) {
    const bridgeId = message.data?.bridgeId || uuidv4();
    setBridgeId(bridgeId);
    bridgeManager.registerBridge(bridgeId, ws, message.data);
    const capacity = bridgeManager.getAvailableCapacity();
    await taskScheduler.setCurrentCapacity(capacity.online);

    const response = {
      type: 'REGISTER_ACK',
      data: {
        bridgeId,
        timestamp: new Date().toISOString()
      }
    };
    ws.send(JSON.stringify(response));
  }

  handleHeartbeat(message) {
    const bridgeId = message.data?.bridgeId;
    if (bridgeId) {
      bridgeManager.updateHeartbeat(bridgeId);
    }
  }

  handleTaskAck(message) {
    const requestId = message.data?.requestId;
    const instanceId = message.data?.instanceId;
    const bridgeId = message.data?.bridgeId;
    
    if (requestId) {
      taskScheduler.handleTaskAck(requestId, instanceId, bridgeId);
      
      if (instanceId) {
        bridgeManager.confirmInstanceLock(instanceId);
      }
      
      if (this.pendingRequests.has(requestId)) {
        const pending = this.pendingRequests.get(requestId);
        pending.ackReceived = true;
        pending.ackAt = new Date().toISOString();
      }
      
      this.instanceTaskMap.set(instanceId, requestId);
    }
    
    this.handleBridgeResponse(message);
  }

  async handleTaskEnd(message) {
    console.log(`[WebSocketServer] 收到 TASK_END 消息：`, JSON.stringify(message.data, null, 2));
    
    const requestId = message.data?.requestId;
    const instanceId = message.data?.instanceId;
    const result = message.data?.result;
    const error = message.data?.error;
    
    console.log(`[WebSocketServer] 解析参数：requestId=${requestId}, instanceId=${instanceId}`);
    
    if (instanceId) {
      const existsInMap = this.instanceTaskMap.has(instanceId);
      console.log(`[WebSocketServer] 实例 ${instanceId} 在 instanceTaskMap 中：${existsInMap}`);
      
      this.instanceTaskMap.delete(instanceId);
      const released = bridgeManager.releaseInstanceLock(instanceId);
      console.log(`[WebSocketServer] 实例锁释放 ${released ? '成功' : '失败'}: ${instanceId}`);
      
      if (!released) {
        console.warn(`[WebSocketServer] 实例锁释放失败，可能原因：1) 锁已超时自动释放 2) 重复收到 TASK_END 消息`);
      }
    } else {
      console.warn(`[WebSocketServer] TASK_END 消息中缺少 instanceId，无法释放锁`);
    }
    
    if (requestId) {
      if (error) {
        await taskScheduler.handleTaskFailure(requestId, error);
      } else {
        await taskScheduler.handleTaskComplete(requestId, result);
      }
    }
    
    this.handleBridgeResponse(message);
  }

  async handleInstanceStatusUpdate(message) {
    const bridgeId = message.data?.bridgeId;
    const instances = message.data?.instances;
    
    if (instances && Array.isArray(instances)) {
      for (const instance of instances) {
        if (instance.status === 'offline') {
          const affectedTaskId = bridgeManager.handleInstanceOffline(instance.id);
          if (affectedTaskId) {
            taskScheduler.recoverTask(affectedTaskId, 'instance_offline');
          }
        }
      }
      const capacity = bridgeManager.getAvailableCapacity();
      await taskScheduler.setCurrentCapacity(capacity.online);
    }
  }

  handleBridgeResponse(message) {
    const requestId = message.data?.requestId;
    if (requestId && this.pendingRequests.has(requestId)) {
      const pending = this.pendingRequests.get(requestId);
      pending.resolve(message);
      this.pendingRequests.delete(requestId);
    }
  }

  sendTaskToBridge(bridgeId, taskData, requestId) {
    return new Promise((resolve, reject) => {
      const message = {
        type: 'TASK_ASSIGN',
        data: {
          ...taskData,
          requestId
        }
      };
      
      console.log('[分发] WebSocketServer 准备发送消息:', JSON.stringify(message, null, 2));

      const success = bridgeManager.sendToBridge(bridgeId, message);
      if (!success) {
        bridgeManager.releaseInstanceLock(taskData.instanceId);
        reject(new Error('发送任务失败'));
        return;
      }

      const timeout = setTimeout(() => {
        if (this.pendingRequests.has(requestId)) {
          this.pendingRequests.delete(requestId);
          bridgeManager.releaseInstanceLock(taskData.instanceId);
          reject(new Error('任务执行超时'));
        }
      }, this.TASK_TIMEOUT);

      this.pendingRequests.set(requestId, { 
        resolve, 
        reject, 
        timeout, 
        bridgeId,
        instanceId: taskData.instanceId,
        sentAt: new Date().toISOString(),
        ackReceived: false
      });
    });
  }

  sendTaskToInstance(bridgeId, instanceId, taskData, requestId) {
    return new Promise((resolve, reject) => {
      const message = {
        type: 'TASK_ASSIGN',
        data: {
          ...taskData,
          requestId,
          instanceId
        }
      };
      
      // console.log(`[分发] WebSocketServer 发送任务到实例：bridgeId=${bridgeId}, instanceId=${instanceId}, requestId=${requestId}`);

      const success = bridgeManager.sendToBridge(bridgeId, message);
      if (!success) {
        bridgeManager.releaseInstanceLock(instanceId);
        reject(new Error('发送任务到实例失败'));
        return;
      }

      const timeout = setTimeout(() => {
        if (this.pendingRequests.has(requestId)) {
          this.pendingRequests.delete(requestId);
          bridgeManager.releaseInstanceLock(instanceId);
          taskScheduler.handleTaskFailure(requestId, '任务发送超时');
          reject(new Error('任务执行超时'));
        }
      }, this.TASK_TIMEOUT);

      this.pendingRequests.set(requestId, { 
        resolve, 
        reject, 
        timeout, 
        bridgeId,
        instanceId,
        sentAt: new Date().toISOString(),
        ackReceived: false
      });
      
      this.instanceTaskMap.set(instanceId, requestId);
    });
  }

  sendInstanceCheckToBridge(bridgeId, checkType, instanceId, requestId) {
    return new Promise((resolve, reject) => {
      const message = {
        type: 'INSTANCE_CHECK',
        data: {
          checkType,
          instanceId,
          requestId
        }
      };
      
      console.log('[分发] WebSocketServer 准备发送实例检查消息:', JSON.stringify(message, null, 2));

      const success = bridgeManager.sendToBridge(bridgeId, message);
      if (!success) {
        reject(new Error('发送实例检查请求失败'));
        return;
      }

      const timeout = setTimeout(() => {
        if (this.pendingRequests.has(requestId)) {
          this.pendingRequests.delete(requestId);
          reject(new Error('实例检查超时'));
        }
      }, 30000);

      this.pendingRequests.set(requestId, { resolve, reject, timeout, bridgeId });
    });
  }

  cleanupPendingRequests(bridgeId) {
    for (const [requestId, pending] of this.pendingRequests) {
      if (pending.bridgeId === bridgeId) {
        clearTimeout(pending.timeout);
        
        if (pending.instanceId) {
          bridgeManager.releaseInstanceLock(pending.instanceId);
          this.instanceTaskMap.delete(pending.instanceId);
          
          taskScheduler.recoverTask(requestId, 'bridge_disconnect');
        }
        
        pending.reject(new Error('桥接器连接已断开'));
        this.pendingRequests.delete(requestId);
      }
    }
  }

  getStats() {
    return {
      pendingRequests: this.pendingRequests.size,
      instanceTaskMap: this.instanceTaskMap.size,
      pendingRequestIds: Array.from(this.pendingRequests.keys())
    };
  }
}

export default new WebSocketServer();