import { WebSocketMessage } from "../types";

type MessageCallback = (data: WebSocketMessage) => void;
type StatusCallback = () => void;
type ErrorCallback = (error: Event) => void;

interface WebSocketClientOptions {
  onMessage?: MessageCallback;
  onOpen?: StatusCallback;
  onClose?: StatusCallback;
  onError?: ErrorCallback;
}

const INITIAL_RECONNECT_DELAY = 1000;
const MAX_RECONNECT_DELAY = 30000;
const RECONNECT_MULTIPLIER = 2;

export class WebSocketClient {
  private ws: WebSocket | null = null;
  private url: string = "";
  private reconnectDelay: number = INITIAL_RECONNECT_DELAY;
  private reconnectTimer: ReturnType<typeof setTimeout> | null = null;
  private shouldReconnect: boolean = false;
  private callbacks: WebSocketClientOptions = {};

  setCallbacks(callbacks: WebSocketClientOptions) {
    this.callbacks = callbacks;
  }

  connect(url: string) {
    this.url = url;
    this.shouldReconnect = true;
    this.reconnectDelay = INITIAL_RECONNECT_DELAY;
    this.openConnection();
  }

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

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

      this.ws.onopen = () => {
        this.reconnectDelay = INITIAL_RECONNECT_DELAY;
        this.callbacks.onOpen?.();
      };

      this.ws.onmessage = (event) => {
        try {
          const data = JSON.parse(event.data) as WebSocketMessage;
          this.callbacks.onMessage?.(data);
        } catch {
          // Non-JSON message — treat as plain text
          const data: WebSocketMessage = {
            type: "text",
            content: String(event.data),
          };
          this.callbacks.onMessage?.(data);
        }
      };

      this.ws.onclose = () => {
        this.callbacks.onClose?.();
        if (this.shouldReconnect) {
          this.scheduleReconnect();
        }
      };

      this.ws.onerror = (error) => {
        this.callbacks.onError?.(error);
      };
    } catch {
      if (this.shouldReconnect) {
        this.scheduleReconnect();
      }
    }
  }

  private scheduleReconnect() {
    if (this.reconnectTimer) {
      clearTimeout(this.reconnectTimer);
    }
    this.reconnectTimer = setTimeout(() => {
      this.reconnectDelay = Math.min(
        this.reconnectDelay * RECONNECT_MULTIPLIER,
        MAX_RECONNECT_DELAY
      );
      this.openConnection();
    }, this.reconnectDelay);
  }

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

  send(message: WebSocketMessage) {
    if (this.ws && this.ws.readyState === WebSocket.OPEN) {
      this.ws.send(JSON.stringify(message));
      return true;
    }
    return false;
  }

  get readyState(): number {
    return this.ws?.readyState ?? WebSocket.CLOSED;
  }

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

export const wsClient = new WebSocketClient();