// ComfyUI WebSocket 实时连接
export class ComfyUIWebSocket {
  constructor(url = 'ws://localhost:8188/ws') {
    this.url = url;
    this.ws = null;
    this.callbacks = {};
    this.isConnected = false;
  }

  connect() {
    return new Promise((resolve, reject) => {
      try {
        this.ws = new WebSocket(this.url);
        
        this.ws.onopen = () => {
          this.isConnected = true;
          console.log('ComfyUI WebSocket connected');
          resolve();
        };

        this.ws.onmessage = (event) => {
          const data = JSON.parse(event.data);
          this.handleMessage(data);
        };

        this.ws.onclose = () => {
          this.isConnected = false;
          console.log('ComfyUI WebSocket disconnected');
        };

        this.ws.onerror = (error) => {
          console.error('ComfyUI WebSocket error:', error);
          reject(error);
        };

      } catch (error) {
        reject(error);
      }
    });
  }

  handleMessage(data) {
    const { type, data: messageData } = data;
    
    switch (type) {
      case 'status':
        this.onStatus(messageData);
        break;
      case 'progress':
        this.onProgress(messageData);
        break;
      case 'executing':
        this.onExecuting(messageData);
        break;
      case 'executed':
        this.onExecuted(messageData);
        break;
      case 'execution_error':
        this.onError(messageData);
        break;
    }
  }

  onStatus(data) {
    if (this.callbacks.onStatus) {
      this.callbacks.onStatus(data);
    }
  }

  onProgress(data) {
    if (this.callbacks.onProgress) {
      this.callbacks.onProgress(data);
    }
  }

  onExecuting(data) {
    if (this.callbacks.onExecuting) {
      this.callbacks.onExecuting(data);
    }
  }

  onExecuted(data) {
    if (this.callbacks.onExecuted) {
      this.callbacks.onExecuted(data);
    }
  }

  onError(data) {
    if (this.callbacks.onError) {
      this.callbacks.onError(data);
    }
  }

  // 设置回调函数
  setCallbacks(callbacks) {
    this.callbacks = callbacks;
  }

  // 发送消息到ComfyUI
  send(message) {
    if (this.isConnected && this.ws) {
      this.ws.send(JSON.stringify(message));
    }
  }

  // 断开连接
  disconnect() {
    if (this.ws) {
      this.ws.close();
    }
  }
}