import type {
  WSMessage,
  WSScanProgress,
  WSScanComplete,
  WSHostDiscovered,
  WSError,
} from '../types/api';

const WS_BASE_URL = import.meta.env.VITE_WS_URL || 'ws://localhost:8000';

export type WSMessageHandler = {
  onScanProgress?: (data: WSScanProgress) => void;
  onScanComplete?: (data: WSScanComplete) => void;
  onHostDiscovered?: (data: WSHostDiscovered) => void;
  onError?: (data: WSError) => void;
  onConnect?: () => void;
  onDisconnect?: () => void;
};

export class WebSocketClient {
  private ws: WebSocket | null = null;
  private handlers: WSMessageHandler = {};
  private reconnectAttempts = 0;
  private maxReconnectAttempts = 5;
  private reconnectDelay = 2000;
  private reconnectTimer: number | null = null;

  constructor(handlers: WSMessageHandler) {
    this.handlers = handlers;
  }

  connect(): void {
    if (this.ws?.readyState === WebSocket.OPEN) {
      return;
    }

    try {
      this.ws = new WebSocket(`${WS_BASE_URL}/api/ws`);

      this.ws.onopen = () => {
        console.log('WebSocket connected');
        this.reconnectAttempts = 0;
        this.handlers.onConnect?.();
      };

      this.ws.onmessage = (event) => {
        try {
          const message: WSMessage = JSON.parse(event.data);
          this.handleMessage(message);
        } catch (error) {
          console.error('Failed to parse WebSocket message:', error);
        }
      };

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

      this.ws.onclose = () => {
        console.log('WebSocket disconnected');
        this.handlers.onDisconnect?.();
        this.attemptReconnect();
      };
    } catch (error) {
      console.error('Failed to create WebSocket connection:', error);
      this.attemptReconnect();
    }
  }

  private handleMessage(message: WSMessage): void {
    switch (message.type) {
      case 'scan_progress':
        this.handlers.onScanProgress?.(message.data as WSScanProgress);
        break;
      case 'scan_complete':
        this.handlers.onScanComplete?.(message.data as WSScanComplete);
        break;
      case 'host_discovered':
        this.handlers.onHostDiscovered?.(message.data as WSHostDiscovered);
        break;
      case 'error':
        this.handlers.onError?.(message.data as WSError);
        break;
      default:
        console.warn('Unknown message type:', message.type);
    }
  }

  private attemptReconnect(): void {
    if (this.reconnectAttempts >= this.maxReconnectAttempts) {
      console.error('Max reconnection attempts reached');
      return;
    }

    if (this.reconnectTimer) {
      return;
    }

    this.reconnectAttempts++;
    const delay = this.reconnectDelay * this.reconnectAttempts;

    console.log(`Attempting to reconnect in ${delay}ms (attempt ${this.reconnectAttempts}/${this.maxReconnectAttempts})`);

    this.reconnectTimer = window.setTimeout(() => {
      this.reconnectTimer = null;
      this.connect();
    }, delay);
  }

  disconnect(): void {
    if (this.reconnectTimer) {
      clearTimeout(this.reconnectTimer);
      this.reconnectTimer = null;
    }

    if (this.ws) {
      this.ws.close();
      this.ws = null;
    }
  }

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