import EventBus from "@/utils/EventBus";
import { WS_ERROR_EVENT } from "@/config/constants";
import { getAuthToken } from "@/utils/auth";

const WS_BASE_URL = import.meta.env.VITE_WS_BASE_URL;

class WebSocketClient {
  private enableAutoReconnect: boolean = true;
  public reconnectAttempts: number = 0;
  private wsInstance: WebSocket | undefined;
  private heartbeatInterval: number = 30;
  private reconnectTimer: NodeJS.Timeout | null = null;
  private heartbeatTimer: NodeJS.Timeout | null = null;
  private reconnectDelay: number = 10;
  private maxReconnectAttempts: number = 10;
  private missedHeartbeats: number = 0;
  private shouldBroadcastError: boolean = false;
  private pendingEventCallbacks: string[] = [];

  constructor(private wsPath: string = "", private protocols: string[] = []) {
    this.initializeConnection();
  }

  private initializeConnection(): void {
    if (this.wsInstance) {
      this.wsInstance.close(1000, "Reinitializing connection");
      this.wsInstance = undefined;
    }

    if (typeof window.WebSocket !== "function") {
      console.error("WebSocket is not supported by your browser");
      return;
    }

    const authToken = getAuthToken();
    const fullUrl = `${WS_BASE_URL}${this.wsPath}${authToken ? `?token=${authToken}` : ""}`;
    this.wsInstance = new WebSocket(fullUrl, this.protocols);

    this.wsInstance.onopen = () => this.handleOpen();
    this.wsInstance.onclose = () => this.handleClose();
    this.wsInstance.onerror = () => this.handleError();
    this.wsInstance.onmessage = (event) => this.handleIncomingMessage(event);
  }

  private handleOpen(): void {
    console.log("WebSocket connection established");
    this.reconnectAttempts = 0;
    if (this.enableAutoReconnect) this.startHeartbeat();
  }

  private handleError(): void {
    console.error("WebSocket connection error occurred");
    this.terminateConnection();
  }

  private handleClose(): void {
    console.log("WebSocket connection closed");
    this.terminateConnection();
  }

  private startHeartbeat(): void {
    if (this.heartbeatTimer) {
      clearInterval(this.heartbeatTimer);
    }

    this.heartbeatTimer = setInterval(() => {
      if (this.isConnectionOpen()) {
        this.send("HEARTBEAT", { type: "PING" });
        this.missedHeartbeats++;

        if (this.missedHeartbeats > 5) {
          clearInterval(this.heartbeatTimer!);
          EventBus.emit(WS_ERROR_EVENT, "Heartbeat timeout: server unresponsive");
          this.terminateConnection();
        }
      } else {
        this.terminateConnection();
      }
    }, this.heartbeatInterval * 1000);
  }

  private formatMessage(eventType: string, payload: any): string {
    const requestId = `${eventType}_${Math.floor(Math.random() * 10000)}`;
    return JSON.stringify({
      event: eventType,
      payload: {
        requestId,
        ...payload
      }
    });
  }

  public send<T = any>(eventType: string, payload?: T, errorHandler?: (msg: string) => void): Promise<{ event: string; requestId: string }> {
    return new Promise((resolve, reject) => {
      if (!this.isConnectionOpen()) {
        console.error("WebSocket is not in connected state");
        this.terminateConnection();
        reject();
        return;
      }

      const message = this.formatMessage(eventType, payload || {});
      const parsedMsg = JSON.parse(message);
      this.wsInstance?.send(message);

      if (eventType !== "HEARTBEAT") {
        console.log("Sent WebSocket message:", message);
      } else {
        this.missedHeartbeats++;
      }

      if (errorHandler && typeof errorHandler === "function") {
        if (!EventBus.hasListener(parsedMsg.payload.requestId)) {
          EventBus.on(parsedMsg.payload.requestId, errorHandler);
          if (!this.pendingEventCallbacks.includes(parsedMsg.payload.requestId)) {
            this.pendingEventCallbacks.push(parsedMsg.payload.requestId);
          }
        }
      } else if (errorHandler) {
        console.error("Third argument to send() must be an error handler function");
      }

      resolve({
        event: eventType,
        requestId: parsedMsg.payload.requestId
      });
    });
  }

  public isConnectionOpen(): boolean {
    const isOpen = this.wsInstance?.readyState === WebSocket.OPEN;
    if (isOpen) this.reconnectAttempts = 0;
    return isOpen;
  }

  private terminateConnection(): void {
    if (this.heartbeatTimer) {
      clearInterval(this.heartbeatTimer);
      this.heartbeatTimer = null;
    }

    if (this.enableAutoReconnect && this.reconnectAttempts < this.maxReconnectAttempts) {
      this.scheduleReconnect();
    } else if (this.reconnectAttempts >= this.maxReconnectAttempts && !this.shouldBroadcastError) {
      console.error(`Failed to reconnect after ${this.maxReconnectAttempts} attempts`);
      EventBus.emit(WS_ERROR_EVENT, "Connection lost permanently");
      this.shouldBroadcastError = true;
    }
  }

  private scheduleReconnect(): void {
    if (!this.reconnectTimer) {
      this.reconnectTimer = setTimeout(() => {
        this.reconnectAttempts++;
        console.log(`Attempting to reconnect (${this.reconnectAttempts}/${this.maxReconnectAttempts})`);
        this.initializeConnection();
        clearTimeout(this.reconnectTimer!);
        this.reconnectTimer = null;
      }, this.reconnectDelay * 1000);
    }
  }

  public disconnect(): void {
    this.enableAutoReconnect = false;
    this.reconnectAttempts = this.maxReconnectAttempts + 1;
    if (this.isConnectionOpen()) {
      this.wsInstance?.close(1000, "Client initiated disconnect");
    }
    this.terminateConnection();
    console.log("WebSocket connection closed by client");
  }

  private handleIncomingMessage(event: MessageEvent): void {
    const { code, payload, requestId, event: eventType, message } = JSON.parse(event.data);

    if (eventType === "HEARTBEAT") {
      this.missedHeartbeats = 0;
      this.reconnectAttempts = 0;
      return;
    }

    if (code === 200) {
      this.cleanupRequestHandler(requestId);
      EventBus.emit(eventType, payload);
      return payload;
    } else {
      EventBus.emit(requestId, message || "Request failed");
      this.cleanupRequestHandler(requestId);
      return null;
    }
  }

  private cleanupRequestHandler(requestId: string, delay: boolean = true): void {
    const timer = setTimeout(() => {
      if (EventBus.hasListener(requestId) && this.pendingEventCallbacks.includes(requestId)) {
        EventBus.off(requestId);
        this.pendingEventCallbacks = this.pendingEventCallbacks.filter(id => id !== requestId);
      }
      clearTimeout(timer);
    }, delay ? 300 : 0);
  }
}

export default WebSocketClient;

Vue 3 Integration

1. Register as Global Property

import { createApp } from 'vue';
import App from './App.vue';
import WebSocketClient from '@/utils/WebSocketClient';

const app = createApp(App);
const wsClient = new WebSocketClient();
app.config.globalProperties.$wsClient = wsClient;

app.mount('#app');

2. Use in Components

<script setup lang="ts">
import { getCurrentInstance, reactive } from 'vue';
import EventBus from '@/utils/EventBus';

const { proxy } = getCurrentInstance() as any;
const wsClient = proxy.$wsClient;

interface Message {
  id: string;
  sender: string;
  content: string;
  timestamp: number;
}

const chatState = reactive({
  currentUserId: 'user_123',
  messages: [] as Message[]
});

// Send a user message
const SEND_USER_MSG = 'SEND_USER_MESSAGE';
wsClient.send(SEND_USER_MSG, {
  recipientId: 'user_456',
  content: 'Hello from Vue 3!',
  type: 'text'
});

// Listen for incoming user messages
EventBus.on(SEND_USER_MSG, (data: Message) => {
  chatState.messages.push(data);
});

// Send system message with error handling
const SEND_SYSTEM_ALERT = 'SEND_SYSTEM_ALERT';
wsClient.send(SEND_SYSTEM_ALERT, {
  type: 'payment',
  amount: 19.99
}, (errorMsg: string) => {
  console.error('System alert failed:', errorMsg);
  // Display error to user
});

// Listen for system message responses
EventBus.on(SEND_SYSTEM_ALERT, (data: any) => {
  console.log('System alert confirmed:', data);
});
</script>

3. Manual Disconnection

import { getCurrentInstance } from 'vue';

const { proxy } = getCurrentInstance() as any;
proxy.$wsClient.disconnect();