ultisuite-client/lib/api/ws.ts

"use client"

import { useEffect } from "react"
import { useQueryClient, type QueryClient } from "@tanstack/react-query"
import type { WsEvent } from "./types"
import { useAuthStore } from "./auth-store"

class WebSocketManager {
  private ws: WebSocket | null = null
  private reconnectAttempts = 0
  private maxReconnectDelay = 30_000
  private reconnectTimer: ReturnType<typeof setTimeout> | null = null
  private lastSeq = 0
  private queryClient: QueryClient | null = null

  init(queryClient: QueryClient) {
    this.queryClient = queryClient
    this.loadLastSeq()
  }

  connect(token: string) {
    if (this.ws?.readyState === WebSocket.OPEN) return

    const baseUrl =
      process.env.NEXT_PUBLIC_WS_URL ??
      (typeof window !== "undefined"
        ? `${window.location.protocol === "https:" ? "wss:" : "ws:"}//${window.location.host}/ws`
        : "")

    const url = `${baseUrl}?token=${encodeURIComponent(token)}&since=${this.lastSeq}`
    this.ws = new WebSocket(url)

    this.ws.onopen = () => {
      this.reconnectAttempts = 0
    }
    this.ws.onmessage = (event) => this.handleMessage(event)
    this.ws.onclose = () => this.scheduleReconnect(token)
    this.ws.onerror = () => {}
  }

  disconnect() {
    if (this.reconnectTimer) clearTimeout(this.reconnectTimer)
    this.ws?.close()
    this.ws = null
  }

  private handleMessage(event: MessageEvent) {
    try {
      const evt: WsEvent = JSON.parse(event.data)
      if (evt.seq) {
        this.lastSeq = evt.seq
        this.saveLastSeq()
      }
      this.handleEvent(evt)
    } catch {}
  }

  private handleEvent(evt: WsEvent) {
    if (!this.queryClient) return

    switch (evt.type) {
      case "mail.created":
        this.queryClient.invalidateQueries({ queryKey: ["messages"] })
        break
      case "mail.updated":
        this.queryClient.invalidateQueries({ queryKey: ["messages"] })
        if (evt.message_id) {
          this.queryClient.invalidateQueries({
            queryKey: ["message", evt.message_id],
          })
        }
        break
      case "mail.deleted":
        this.queryClient.invalidateQueries({ queryKey: ["messages"] })
        if (evt.message_id) {
          this.queryClient.removeQueries({
            queryKey: ["message", evt.message_id],
          })
        }
        break
      case "outbox.updated":
        this.queryClient.invalidateQueries({ queryKey: ["outbox"] })
        break
      case "contact.updated":
        this.queryClient.invalidateQueries({ queryKey: ["contacts"] })
        break
    }
  }

  private scheduleReconnect(token: string) {
    const delay = Math.min(
      1000 * 2 ** this.reconnectAttempts,
      this.maxReconnectDelay
    )
    this.reconnectAttempts++
    this.reconnectTimer = setTimeout(() => this.connect(token), delay)
  }

  private loadLastSeq() {
    if (typeof window === "undefined") return
    const stored = localStorage.getItem("ultimail-ws-seq")
    if (stored) this.lastSeq = parseInt(stored, 10) || 0
  }

  private saveLastSeq() {
    if (typeof window === "undefined") return
    localStorage.setItem("ultimail-ws-seq", String(this.lastSeq))
  }
}

export const wsManager = new WebSocketManager()

export function useWebSocket() {
  const queryClient = useQueryClient()
  const accessToken = useAuthStore((s) => s.accessToken)

  useEffect(() => {
    wsManager.init(queryClient)
  }, [queryClient])

  useEffect(() => {
    if (accessToken) {
      wsManager.connect(accessToken)
    } else {
      wsManager.disconnect()
    }
    return () => wsManager.disconnect()
  }, [accessToken])
}