type ProgressCallback = (data: ProgressMessage) => void;

interface ProgressMessage {
  type: 'progress' | 'status' | 'error' | 'complete';
  taskId?: string;
  filename?: string;
  progress?: number;
  status?: string;
  message?: string;
  timestamp?: string;
}

class ProgressWebSocket {
  private ws: WebSocket | null = null;
  private url: string;
  private callbacks: Set<ProgressCallback> = new Set();
  private reconnectTimer: ReturnType<typeof setTimeout> | null = null;
  private reconnectInterval = 3000;
  private maxReconnectInterval = 30000;
  private shouldReconnect = false;
  private shouldCloseAfterOpen = false;
  private currentInterval = 3000;

  constructor(url = 'ws://192.168.3.11:8000/ws/progress') {
    this.url = url;
  }

  connect() {
    if (this.ws && (this.ws.readyState === WebSocket.OPEN || this.ws.readyState === WebSocket.CONNECTING)) {
      return;
    }

    this.shouldReconnect = true;
    this.shouldCloseAfterOpen = false;

    try {
      this.ws = new WebSocket(this.url);

      this.ws.onopen = () => {
        // 如果连接前被要求关闭（React StrictMode 快速卸载），则打开后立即关闭
        if (this.shouldCloseAfterOpen) {
          this.ws?.close();
          this.ws = null;
          return;
        }
        this.currentInterval = this.reconnectInterval;
        console.log('[WebSocket] Connected to progress stream');
      };

      this.ws.onmessage = (event) => {
        try {
          const data: ProgressMessage = JSON.parse(event.data);
          this.callbacks.forEach((cb) => cb(data));
        } catch (err) {
          console.error('[WebSocket] Failed to parse message:', err);
        }
      };

      this.ws.onclose = () => {
        console.log('[WebSocket] Connection closed');
        this.ws = null;
        if (this.shouldReconnect) {
          this.scheduleReconnect();
        }
      };

      this.ws.onerror = () => {
        // 静默处理错误，避免在 CONNECTING 状态时 close 触发浏览器报错
        this.ws = null;
        if (this.shouldReconnect) {
          this.scheduleReconnect();
        }
      };
    } catch (err) {
      console.error('[WebSocket] Failed to connect:', err);
      this.scheduleReconnect();
    }
  }

  disconnect() {
    this.shouldReconnect = false;
    if (this.reconnectTimer) {
      clearTimeout(this.reconnectTimer);
      this.reconnectTimer = null;
    }
    if (!this.ws) return;

    // 如果还在连接中，设置标志位让 onopen 打开后立即关闭
    if (this.ws.readyState === WebSocket.CONNECTING) {
      this.shouldCloseAfterOpen = true;
      return;
    }

    // 只有已打开的连接才调用 close()
    if (this.ws.readyState === WebSocket.OPEN) {
      this.ws.close();
    }
    this.ws = null;
  }

  onProgress(callback: ProgressCallback) {
    this.callbacks.add(callback);
    return () => {
      this.callbacks.delete(callback);
    };
  }

  private scheduleReconnect() {
    if (this.reconnectTimer) {
      clearTimeout(this.reconnectTimer);
    }
    this.reconnectTimer = setTimeout(() => {
      console.log(`[WebSocket] Reconnecting in ${this.currentInterval}ms...`);
      this.connect();
      this.currentInterval = Math.min(this.currentInterval * 1.5, this.maxReconnectInterval);
    }, this.currentInterval);
  }

  isConnected(): boolean {
    return this.ws !== null && this.ws.readyState === WebSocket.OPEN;
  }
}

export const progressWS = new ProgressWebSocket();
export type { ProgressMessage };
