import { WS_PROGRESS_URL } from './config';

type ProgressCallback = (data: ProgressMessage) => void;
type ConnectionStatus = 'connecting' | 'connected' | 'reconnecting' | 'disconnected';
type StatusCallback = (status: ConnectionStatus) => void;

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

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

  constructor(url = WS_PROGRESS_URL) {
    this.url = url;
  }

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

    this.shouldReconnect = true;
    this.shouldCloseAfterOpen = false;
    this.notifyStatus('connecting');

    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;
        this.startHeartbeat();
        this.notifyStatus('connected');
        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.stopHeartbeat();
        this.ws = null;
        this.notifyStatus('disconnected');
        if (this.shouldReconnect) {
          this.scheduleReconnect();
        }
      };

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

  disconnect() {
    this.shouldReconnect = false;
    this.stopHeartbeat();
    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;
    this.notifyStatus('disconnected');
  }

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

  onStatus(callback: StatusCallback) {
    this.statusCallbacks.add(callback);
    callback(this.isConnected() ? 'connected' : 'disconnected');
    return () => {
      this.statusCallbacks.delete(callback);
    };
  }

  private scheduleReconnect() {
    if (this.reconnectTimer) {
      clearTimeout(this.reconnectTimer);
    }
    this.notifyStatus('reconnecting');
    this.reconnectTimer = setTimeout(() => {
      console.log('[WebSocket] Reconnecting to progress stream...');
      this.connect();
      this.currentInterval = Math.min(this.currentInterval * 1.5, this.maxReconnectInterval);
    }, this.currentInterval);
  }

  private startHeartbeat() {
    this.stopHeartbeat();
    this.sendHeartbeat();
    this.heartbeatTimer = setInterval(() => this.sendHeartbeat(), this.heartbeatInterval);
  }

  private stopHeartbeat() {
    if (this.heartbeatTimer) {
      clearInterval(this.heartbeatTimer);
      this.heartbeatTimer = null;
    }
  }

  private sendHeartbeat() {
    if (this.ws?.readyState === WebSocket.OPEN) {
      this.ws.send('ping');
    }
  }

  private notifyStatus(status: ConnectionStatus) {
    this.statusCallbacks.forEach((cb) => cb(status));
  }

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

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